Upload 21 files

ag4masses/alphageometry/alphageometry.py

@@ -33,8 +33,6 @@ import problem as pr
 #=============
 import sys, os, math, re
 import multiprocessing
-import warnings
-warnings.filterwarnings("ignore")
 model = None # global variable used in multi-processing workers
 
 _GIN_SEARCH_PATHS = flags.DEFINE_list(
@@ -152,8 +150,8 @@ def write_solution(g: gh.Graph, p: pr.Problem, out_file: str) -> None:
       g, p.goal, merge_trivials=False
   )
 
-  solution = ''
-  solution += 'Theo đề bài ta có:\n'
+  solution = '\n=========================='
+  solution += '\n * From theorem premises:\n'
   premises_nl = []
   for premises, [points] in setup:
     solution += ' '.join([p.name.upper() for p in points]) + ' '
@@ -165,18 +163,15 @@ def write_solution(g: gh.Graph, p: pr.Problem, out_file: str) -> None:
     ]
   solution += ': Points\n' + '\n'.join(premises_nl)
 
-  solution += '\n\nCác điểm cần dựng thêm:\n'
+  solution += '\n\n * Auxiliary Constructions:\n'
   aux_premises_nl = []
-  if len(aux) == 0:
-    solution += 'Không cần dựng thêm điểm nào.'
-  else:
-    for premises, [points] in aux:
-      solution += ' '.join([p.name.upper() for p in points]) + ' '
-      aux_premises_nl += [
-          natural_language_statement(p) + ' [{:02}]'.format(refs[p.hashed()])
-          for p in premises
-      ]
-    solution += ': Points\n' + '\n'.join(aux_premises_nl)
+  for premises, [points] in aux:
+    solution += ' '.join([p.name.upper() for p in points]) + ' '
+    aux_premises_nl += [
+        natural_language_statement(p) + ' [{:02}]'.format(refs[p.hashed()])
+        for p in premises
+    ]
+  solution += ': Points\n' + '\n'.join(aux_premises_nl)
 
   # some special case where the deduction rule has a well known name.
   r2name = {
@@ -194,19 +189,22 @@ def write_solution(g: gh.Graph, p: pr.Problem, out_file: str) -> None:
       'a02': '(Angle chase)',
   }
 
-  solution += '\n\nCác bước chứng minh:\n'
+  solution += '\n\n * Proof steps:\n'
   for i, step in enumerate(proof_steps):
     _, [con] = step
     nl = proof_step_string(step, refs, last_step=i == len(proof_steps) - 1)
     rule_name = r2name.get(con.rule_name, '')
     nl = nl.replace('\u21d2', f'{rule_name}\u21d2 ')
     solution += '{:03}. '.format(i + 1) + nl + '\n'
+
+  solution += '==========================\n'
   logging.info(solution)
   if out_file:
     with open(out_file, 'w') as f:
       f.write(solution)
     logging.info('Solution written to %s.', out_file)
 
+
 def get_lm(ckpt_init: str, vocab_path: str) -> lm.LanguageModelInference:
   lm.parse_gin_configuration(
       _GIN_FILE.value, _GIN_PARAM.value, gin_paths=_GIN_SEARCH_PATHS.value
@@ -234,10 +232,12 @@ def run_ddar(g: gh.Graph, p: pr.Problem, out_file: str) -> bool:
     return False
 
   write_solution(g, p, out_file)
+
   gh.nm.draw(
       g.type2nodes[gh.Point],
       g.type2nodes[gh.Line],
       g.type2nodes[gh.Circle],
+      g.type2nodes[gh.SemiCircle],
       g.type2nodes[gh.Segment],
       goal=(p.goal.name, goal_args),
       save_to="ag4mout/output.png",)
@@ -718,7 +718,7 @@ def main(_):
   # point names will be renamed to alphabetical a, b, c, d, e, ...
   # instead of staying with their original names,
   # in order to match the synthetic training data generation.
-  need_rename = _MODE.value != 'ddar'
+  need_rename = False
 
   # load problems from the problems_file,
   problems = pr.Problem.from_txt_file(
@@ -752,4 +752,4 @@ def main(_):
 
 
 if __name__ == '__main__':
-  app.run(main)
+  app.run(main)

ag4masses/alphageometry/dd.py

@@ -1,1156 +1,1220 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Implements Deductive Database (DD)."""
-
-# pylint: disable=g-multiple-import,g-importing-member
-from collections import defaultdict
-import time
-from typing import Any, Callable, Generator
-
-import geometry as gm
-import graph as gh
-import graph_utils as utils
-import numericals as nm
-import problem as pr
-from problem import Dependency, EmptyDependency
-
-
-def intersect1(set1: set[Any], set2: set[Any]) -> Any:
-  for x in set1:
-    if x in set2:
-      return x
-  return None
-
-
-def diff_point(l: gm.Line, a: gm.Point) -> gm.Point:
-  for x in l.neighbors(gm.Point):
-    if x != a:
-      return x
-  return None
-
-
-# pylint: disable=protected-access
-# pylint: disable=unused-argument
-
-
-def match_eqratio_eqratio_eqratio(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqratio a b c d m n p q, eqratio c d e f p q r u => eqratio a b e f m n r u."""
-  for m1 in g.type2nodes[gm.Value]:
-    for m2 in g.type2nodes[gm.Value]:
-      rats1 = []
-      for rat in m1.neighbors(gm.Ratio):
-        l1, l2 = rat.lengths
-        if l1 is None or l2 is None:
-          continue
-        rats1.append((l1, l2))
-
-      rats2 = []
-      for rat in m2.neighbors(gm.Ratio):
-        l1, l2 = rat.lengths
-        if l1 is None or l2 is None:
-          continue
-        rats2.append((l1, l2))
-
-      pairs = []
-      for (l1, l2), (l3, l4) in utils.cross(rats1, rats2):
-        if l2 == l3:
-          pairs.append((l1, l2, l4))
-
-      for (l1, l12, l2), (l3, l34, l4) in utils.comb2(pairs):
-        if (l1, l12, l2) == (l3, l34, l4):
-          continue
-        if l1 == l2 or l3 == l4:
-          continue
-        if l1 == l12 or l12 == l2 or l3 == l34 or l4 == l34:
-          continue
-        # d12 - d1 = d34 - d3 = m1
-        # d2 - d12 = d4 - d34 = m2
-        # => d2 - d1 = d4 - d3 (= m1+m2)
-        a, b = g.two_points_of_length(l1)
-        c, d = g.two_points_of_length(l12)
-        m, n = g.two_points_of_length(l3)
-        p, q = g.two_points_of_length(l34)
-        # eqangle a b c d m n p q
-        e, f = g.two_points_of_length(l2)
-        r, u = g.two_points_of_length(l4)
-        yield dict(zip('abcdefmnpqru', [a, b, c, d, e, f, m, n, p, q, r, u]))
-
-
-def match_eqangle_eqangle_eqangle(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle a b c d m n p q, eqangle c d e f p q r u => eqangle a b e f m n r u."""
-  for m1 in g.type2nodes[gm.Measure]:
-    for m2 in g.type2nodes[gm.Measure]:
-      angs1 = []
-      for ang in m1.neighbors(gm.Angle):
-        d1, d2 = ang.directions
-        if d1 is None or d2 is None:
-          continue
-        angs1.append((d1, d2))
-
-      angs2 = []
-      for ang in m2.neighbors(gm.Angle):
-        d1, d2 = ang.directions
-        if d1 is None or d2 is None:
-          continue
-        angs2.append((d1, d2))
-
-      pairs = []
-      for (d1, d2), (d3, d4) in utils.cross(angs1, angs2):
-        if d2 == d3:
-          pairs.append((d1, d2, d4))
-
-      for (d1, d12, d2), (d3, d34, d4) in utils.comb2(pairs):
-        if (d1, d12, d2) == (d3, d34, d4):
-          continue
-        if d1 == d2 or d3 == d4:
-          continue
-        if d1 == d12 or d12 == d2 or d3 == d34 or d4 == d34:
-          continue
-        # d12 - d1 = d34 - d3 = m1
-        # d2 - d12 = d4 - d34 = m2
-        # => d2 - d1 = d4 - d3
-        a, b = g.two_points_on_direction(d1)
-        c, d = g.two_points_on_direction(d12)
-        m, n = g.two_points_on_direction(d3)
-        p, q = g.two_points_on_direction(d34)
-        # eqangle a b c d m n p q
-        e, f = g.two_points_on_direction(d2)
-        r, u = g.two_points_on_direction(d4)
-        yield dict(zip('abcdefmnpqru', [a, b, c, d, e, f, m, n, p, q, r, u]))
-
-
-def match_perp_perp_npara_eqangle(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match perp A B C D, perp E F G H, npara A B E F => eqangle A B E F C D G H."""
-  dpairs = []
-  for ang in g.vhalfpi.neighbors(gm.Angle):
-    d1, d2 = ang.directions
-    if d1 is None or d2 is None:
-      continue
-    dpairs.append((d1, d2))
-
-  for (d1, d2), (d3, d4) in utils.comb2(dpairs):
-    a, b = g.two_points_on_direction(d1)
-    c, d = g.two_points_on_direction(d2)
-    m, n = g.two_points_on_direction(d3)
-    p, q = g.two_points_on_direction(d4)
-    if g.check_npara([a, b, m, n]):
-      if ({a, b}, {c, d}) == ({m, n}, {p, q}):
-        continue
-      if ({a, b}, {c, d}) == ({p, q}, {m, n}):
-        continue
-
-      yield dict(zip('ABCDEFGH', [a, b, c, d, m, n, p, q]))
-
-
-def match_circle_coll_eqangle_midp(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match circle O A B C, coll M B C, eqangle A B A C O B O M => midp M B C."""
-  for p, a, b, c in g.all_circles():
-    ab = g._get_line(a, b)
-    if ab is None:
-      continue
-    if ab.val is None:
-      continue
-    ac = g._get_line(a, c)
-    if ac is None:
-      continue
-    if ac.val is None:
-      continue
-    pb = g._get_line(p, b)
-    if pb is None:
-      continue
-    if pb.val is None:
-      continue
-
-    bc = g._get_line(b, c)
-    if bc is None:
-      continue
-    bc_points = bc.neighbors(gm.Point, return_set=True)
-
-    anga, _ = g._get_angle(ab.val, ac.val)
-
-    for angp in pb.val.neighbors(gm.Angle):
-      if not g.is_equal(anga, angp):
-        continue
-
-      _, d = angp.directions
-      for l in d.neighbors(gm.Line):
-        l_points = l.neighbors(gm.Point, return_set=True)
-        m = intersect1(bc_points, l_points)
-        if m is not None:
-          yield dict(zip('ABCMO', [a, b, c, m, p]))
-
-
-def match_midp_perp_cong(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match midp M A B, perp O M A B => cong O A O B."""
-  for m, a, b in g.all_midps():
-    ab = g._get_line(a, b)
-    for l in m.neighbors(gm.Line):
-      if g.check_perpl(l, ab):
-        for o in l.neighbors(gm.Point):
-          if o != m:
-            yield dict(zip('ABMO', [a, b, m, o]))
-
-
-def match_cyclic_eqangle_cong(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match cyclic A B C P Q R, eqangle C A C B R P R Q => cong A B P Q."""
-  for c in g.type2nodes[gm.Circle]:
-    ps = c.neighbors(gm.Point)
-    for (a, b, c), (x, y, z) in utils.comb2(list(utils.perm3(ps))):
-      if {a, b, c} == {x, y, z}:
-        continue
-      if g.check_eqangle([c, a, c, b, z, x, z, y]):
-        yield dict(zip('ABCPQR', [a, b, c, x, y, z]))
-
-
-def match_circle_eqangle_perp(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match circle O A B C, eqangle A X A B C A C B => perp O A A X."""
-  for p, a, b, c in g.all_circles():
-    ca = g._get_line(c, a)
-    if ca is None:
-      continue
-    cb = g._get_line(c, b)
-    if cb is None:
-      continue
-    ab = g._get_line(a, b)
-    if ab is None:
-      continue
-
-    if ca.val is None:
-      continue
-    if cb.val is None:
-      continue
-    if ab.val is None:
-      continue
-
-    c_ang, _ = g._get_angle(cb.val, ca.val)
-    if c_ang is None:
-      continue
-
-    for ang in ab.val.neighbors(gm.Angle):
-      if g.is_equal(ang, c_ang):
-        _, d = ang.directions
-        for l in d.neighbors(gm.Line):
-          if a not in l.neighbors(gm.Point):
-            continue
-          x = diff_point(l, a)
-          if x is None:
-            continue
-          yield dict(zip('OABCX', [p, a, b, c, x]))
-        break
-
-
-def match_circle_perp_eqangle(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match circle O A B C, perp O A A X => eqangle A X A B C A C B."""
-  for p, a, b, c in g.all_circles():
-    pa = g._get_line(p, a)
-    if pa is None:
-      continue
-    if pa.val is None:
-      continue
-    for l in a.neighbors(gm.Line):
-      if g.check_perpl(pa, l):
-        x = diff_point(l, a)
-        if x is not None:
-          yield dict(zip('OABCX', [p, a, b, c, x]))
-
-
-def match_perp_perp_ncoll_para(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match perp A B C D, perp C D E F, ncoll A B E => para A B E F."""
-  d2d = defaultdict(list)
-  for ang in g.vhalfpi.neighbors(gm.Angle):
-    d1, d2 = ang.directions
-    if d1 is None or d2 is None:
-      continue
-    d2d[d1] += [d2]
-    d2d[d2] += [d1]
-
-  for x, ys in d2d.items():
-    if len(ys) < 2:
-      continue
-    c, d = g.two_points_on_direction(x)
-    for y1, y2 in utils.comb2(ys):
-      a, b = g.two_points_on_direction(y1)
-      e, f = g.two_points_on_direction(y2)
-      if nm.check_ncoll([a.num, b.num, e.num]):
-        yield dict(zip('ABCDEF', [a, b, c, d, e, f]))
-
-
-def match_eqangle6_ncoll_cong(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 A O A B B A B O, ncoll O A B => cong O A O B."""
-  for a in g.type2nodes[gm.Point]:
-    for b, c in utils.comb2(g.type2nodes[gm.Point]):
-      if a == b or a == c:
-        continue
-      if g.check_eqangle([b, a, b, c, c, b, c, a]):
-        if g.check_ncoll([a, b, c]):
-          yield dict(zip('OAB', [a, b, c]))
-
-
-def match_eqangle_perp_perp(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle A B P Q C D U V, perp P Q U V => perp A B C D."""
-  for ang in g.vhalfpi.neighbors(gm.Angle):
-    # d1 perp d2
-    d1, d2 = ang.directions
-    if d1 is None or d2 is None:
-      continue
-    for d3, d4 in utils.comb2(g.type2nodes[gm.Direction]):
-      if d1 == d3 or d2 == d4:
-        continue
-      # if d1 - d3 = d2 - d4 => d3 perp d4
-      a13, a31 = g._get_angle(d1, d3)
-      a24, a42 = g._get_angle(d2, d4)
-      if a13 is None or a31 is None or a24 is None or a42 is None:
-        continue
-      if g.is_equal(a13, a24) and g.is_equal(a31, a42):
-        a, b = g.two_points_on_direction(d1)
-        c, d = g.two_points_on_direction(d2)
-        m, n = g.two_points_on_direction(d3)
-        p, q = g.two_points_on_direction(d4)
-        yield dict(zip('ABCDPQUV', [m, n, p, q, a, b, c, d]))
-
-
-def match_eqangle_ncoll_cyclic(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 P A P B Q A Q B, ncoll P Q A B => cyclic A B P Q."""
-  for l1, l2, l3, l4 in g.all_eqangles_distinct_linepairss():
-    if len(set([l1, l2, l3, l4])) < 4:
-      continue  # they all must be distinct.
-
-    p1s = l1.neighbors(gm.Point, return_set=True)
-    p2s = l2.neighbors(gm.Point, return_set=True)
-    p3s = l3.neighbors(gm.Point, return_set=True)
-    p4s = l4.neighbors(gm.Point, return_set=True)
-
-    p = intersect1(p1s, p2s)
-    if not p:
-      continue
-    q = intersect1(p3s, p4s)
-    if not q:
-      continue
-    a = intersect1(p1s, p3s)
-    if not a:
-      continue
-    b = intersect1(p2s, p4s)
-    if not b:
-      continue
-    if len(set([a, b, p, q])) < 4:
-      continue
-
-    if not g.check_ncoll([a, b, p, q]):
-      continue
-
-    yield dict(zip('ABPQ', [a, b, p, q]))
-
-
-def match_eqangle_para(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle A B P Q C D P Q => para A B C D."""
-  for measure in g.type2nodes[gm.Measure]:
-    angs = measure.neighbors(gm.Angle)
-    d12, d21 = defaultdict(list), defaultdict(list)
-    for ang in angs:
-      d1, d2 = ang.directions
-      if d1 is None or d2 is None:
-        continue
-      d12[d1].append(d2)
-      d21[d2].append(d1)
-
-    for d1, d2s in d12.items():
-      a, b = g.two_points_on_direction(d1)
-      for d2, d3 in utils.comb2(d2s):
-        c, d = g.two_points_on_direction(d2)
-        e, f = g.two_points_on_direction(d3)
-        yield dict(zip('ABCDPQ', [c, d, e, f, a, b]))
-
-
-def match_cyclic_eqangle(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match cyclic A B P Q => eqangle P A P B Q A Q B."""
-  record = set()
-  for a, b, c, d in g_matcher('cyclic'):
-    if (a, b, c, d) in record:
-      continue
-    record.add((a, b, c, d))
-    record.add((a, b, d, c))
-    record.add((b, a, c, d))
-    record.add((b, a, d, c))
-    yield dict(zip('ABPQ', [a, b, c, d]))
-
-
-def rotate_simtri(
-    a: gm.Point, b: gm.Point, c: gm.Point, x: gm.Point, y: gm.Point, z: gm.Point
-) -> Generator[tuple[gm.Point, ...], None, None]:
-  """Rotate points around for similar triangle predicates."""
-  yield (z, y, x, c, b, a)
-  for p in [
-      (b, c, a, y, z, x),
-      (c, a, b, z, x, y),
-      (x, y, z, a, b, c),
-      (y, z, x, b, c, a),
-      (z, x, y, c, a, b),
-  ]:
-    yield p
-    yield p[::-1]
-
-
-def match_cong_cong_cong_cyclic(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match cong O A O B, cong O B O C, cong O C O D => cyclic A B C D."""
-  for l in g.type2nodes[gm.Length]:
-    p2p = defaultdict(list)
-    for s in l.neighbors(gm.Segment):
-      a, b = s.points
-      p2p[a].append(b)
-      p2p[b].append(a)
-
-    for p, ps in p2p.items():
-      if len(ps) >= 4:
-        for a, b, c, d in utils.comb4(ps):
-          yield dict(zip('OABCD', [p, a, b, c, d]))
-
-
-def match_cong_cong_cong_ncoll_contri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match cong A B P Q, cong B C Q R, cong C A R P, ncoll A B C => contri* A B C P Q R."""
-  record = set()
-  for a, b, p, q in g_matcher('cong'):
-    for c in g.type2nodes[gm.Point]:
-      for r in g.type2nodes[gm.Point]:
-        if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
-          continue
-        if not g.check_ncoll([a, b, c]):
-          continue
-        if g.check_cong([b, c, q, r]) and g.check_cong([c, a, r, p]):
-          record.add((a, b, c, p, q, r))
-          yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
-
-
-def match_cong_cong_eqangle6_ncoll_contri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match cong A B P Q, cong B C Q R, eqangle6 B A B C Q P Q R, ncoll A B C => contri* A B C P Q R."""
-  record = set()
-  for a, b, p, q in g_matcher('cong'):
-    for c in g.type2nodes[gm.Point]:
-      if c in (a, b):
-        continue
-      for r in g.type2nodes[gm.Point]:
-        if r in (p, q):
-          continue
-
-        in_record = False
-        for x in [
-            (c, b, a, r, q, p),
-            (p, q, r, a, b, c),
-            (r, q, p, c, b, a),
-        ]:
-          if x in record:
-            in_record = True
-            break
-
-        if in_record:
-          continue
-
-        if not g.check_cong([b, c, q, r]):
-          continue
-        if not g.check_ncoll([a, b, c]):
-          continue
-
-        if nm.same_clock(a.num, b.num, c.num, p.num, q.num, r.num):
-          if g.check_eqangle([b, a, b, c, q, p, q, r]):
-            record.add((a, b, c, p, q, r))
-            yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
-        else:
-          if g.check_eqangle([b, a, b, c, q, r, q, p]):
-            record.add((a, b, c, p, q, r))
-            yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
-
-
-def match_eqratio6_eqangle6_ncoll_simtri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C => simtri* A B C P Q R."""
-  enums = g_matcher('eqratio6')
-
-  record = set()
-  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    if nm.same_clock(a.num, b.num, c.num, p.num, q.num, r.num):
-      if g.check_eqangle([b, a, b, c, q, p, q, r]):
-        record.add((a, b, c, p, q, r))
-        yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    elif g.check_eqangle([b, a, b, c, q, r, q, p]):
-      record.add((a, b, c, p, q, r))
-      yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
-
-
-def match_eqangle6_eqangle6_ncoll_simtri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 B A B C Q P Q R, eqangle6 C A C B R P R Q, ncoll A B C => simtri A B C P Q R."""
-  enums = g_matcher('eqangle6')
-
-  record = set()
-  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqangle([c, a, c, b, r, p, r, q]):
-      continue
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def match_eqratio6_eqratio6_ncoll_simtri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C => simtri* A B C P Q R."""
-  enums = g_matcher('eqratio6')
-
-  record = set()
-  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqratio([c, a, c, b, r, p, r, q]):
-      continue
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def match_eqangle6_eqangle6_ncoll_simtri2(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 B A B C Q R Q P, eqangle6 C A C B R Q R P, ncoll A B C => simtri2 A B C P Q R."""
-  enums = g_matcher('eqangle6')
-
-  record = set()
-  for b, a, b, c, q, r, q, p in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqangle([c, a, c, b, r, q, r, p]):
-      continue
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def rotate_contri(
-    a: gm.Point, b: gm.Point, c: gm.Point, x: gm.Point, y: gm.Point, z: gm.Point
-) -> Generator[tuple[gm.Point, ...], None, None]:
-  for p in [(b, a, c, y, x, z), (x, y, z, a, b, c), (y, x, z, b, a, c)]:
-    yield p
-
-
-def match_eqangle6_eqangle6_ncoll_cong_contri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 B A B C Q P Q R, eqangle6 C A C B R P R Q, ncoll A B C, cong A B P Q => contri A B C P Q R."""
-  enums = g_matcher('eqangle6')
-
-  record = set()
-  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if not g.check_cong([a, b, p, q]):
-      continue
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqangle([c, a, c, b, r, p, r, q]):
-      continue
-
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def match_eqratio6_eqratio6_ncoll_cong_contri(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C, cong A B P Q => contri* A B C P Q R."""
-  enums = g_matcher('eqratio6')
-
-  record = set()
-  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if not g.check_cong([a, b, p, q]):
-      continue
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqratio([c, a, c, b, r, p, r, q]):
-      continue
-
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def match_eqangle6_eqangle6_ncoll_cong_contri2(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 B A B C Q R Q P, eqangle6 C A C B R Q R P, ncoll A B C, cong A B P Q => contri2 A B C P Q R."""
-  enums = g_matcher('eqangle6')
-
-  record = set()
-  for b, a, b, c, q, r, q, p in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
-    if not g.check_cong([a, b, p, q]):
-      continue
-    if (a, b, c) == (p, q, r):
-      continue
-    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
-      continue
-    if not g.check_eqangle([c, a, c, b, r, q, r, p]):
-      continue
-    if not g.check_ncoll([a, b, c]):
-      continue
-
-    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
-    record.add((a, b, c, p, q, r))
-    yield mapping
-
-
-def match_eqratio6_coll_ncoll_eqangle6(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqratio6 d b d c a b a c, coll d b c, ncoll a b c => eqangle6 a b a d a d a c."""
-  records = set()
-  for b, d, c in g_matcher('coll'):
-    for a in g.all_points():
-      if g.check_coll([a, b, c]):
-        continue
-      if (a, b, d, c) in records or (a, c, d, b) in records:
-        continue
-      records.add((a, b, d, c))
-
-      if g.check_eqratio([d, b, d, c, a, b, a, c]):
-        yield dict(zip('abcd', [a, b, c, d]))
-
-
-def match_eqangle6_coll_ncoll_eqratio6(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 a b a d a d a c, coll d b c, ncoll a b c => eqratio6 d b d c a b a c."""
-  records = set()
-  for b, d, c in g_matcher('coll'):
-    for a in g.all_points():
-      if g.check_coll([a, b, c]):
-        continue
-      if (a, b, d, c) in records or (a, c, d, b) in records:
-        continue
-      records.add((a, b, d, c))
-
-      if g.check_eqangle([a, b, a, d, a, d, a, c]):
-        yield dict(zip('abcd', [a, b, c, d]))
-
-
-def match_eqangle6_ncoll_cyclic(
-    g: gh.Graph,
-    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem,
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match eqangle6 P A P B Q A Q B, ncoll P Q A B => cyclic A B P Q."""
-  for a, b, a, c, x, y, x, z in g_matcher('eqangle6'):  # pylint: disable=redeclared-assigned-name,unused-variable
-    if (b, c) != (y, z) or a == x:
-      continue
-    if nm.check_ncoll([x.num for x in [a, b, c, x]]):
-      yield dict(zip('ABPQ', [b, c, a, x]))
-
-
-def match_all(
-    name: str, g: gh.Graph
-) -> Generator[tuple[gm.Point, ...], None, None]:
-  """Match all instances of a certain relation."""
-  if name in ['ncoll', 'npara', 'nperp']:
-    return []
-  if name == 'coll':
-    return g.all_colls()
-  if name == 'para':
-    return g.all_paras()
-  if name == 'perp':
-    return g.all_perps()
-  if name == 'cong':
-    return g.all_congs()
-  if name == 'eqangle':
-    return g.all_eqangles_8points()
-  if name == 'eqangle6':
-    return g.all_eqangles_6points()
-  if name == 'eqratio':
-    return g.all_eqratios_8points()
-  if name == 'eqratio6':
-    return g.all_eqratios_6points()
-  if name == 'cyclic':
-    return g.all_cyclics()
-  if name == 'midp':
-    return g.all_midps()
-  if name == 'circle':
-    return g.all_circles()
-  raise ValueError(f'Unrecognize {name}')
-
-
-def cache_match(
-    graph: gh.Graph,
-) -> Callable[str, list[tuple[gm.Point, ...]]]:
-  """Cache throughout one single BFS level."""
-  cache = {}
-
-  def match_fn(name: str) -> list[tuple[gm.Point, ...]]:
-    if name in cache:
-      return cache[name]
-
-    result = list(match_all(name, graph))
-    cache[name] = result
-    return result
-
-  return match_fn
-
-
-def try_to_map(
-    clause_enum: list[tuple[pr.Clause, list[tuple[gm.Point, ...]]]],
-    mapping: dict[str, gm.Point],
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Recursively try to match the remaining points given current mapping."""
-  if not clause_enum:
-    yield mapping
-    return
-
-  clause, enum = clause_enum[0]
-  for points in enum:
-    mpcpy = dict(mapping)
-
-    fail = False
-    for p, a in zip(points, clause.args):
-      if a in mpcpy and mpcpy[a] != p or p in mpcpy and mpcpy[p] != a:
-        fail = True
-        break
-      mpcpy[a] = p
-      mpcpy[p] = a
-
-    if fail:
-      continue
-
-    for m in try_to_map(clause_enum[1:], mpcpy):
-      yield m
-
-
-def match_generic(
-    g: gh.Graph,
-    cache: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match any generic rule that is not one of the above match_*() rules."""
-  clause2enum = {}
-
-  clauses = []
-  numerical_checks = []
-  for clause in theorem.premise:
-    if clause.name in ['ncoll', 'npara', 'nperp', 'sameside']:
-      numerical_checks.append(clause)
-      continue
-
-    enum = cache(clause.name)
-    if len(enum) == 0:  # pylint: disable=g-explicit-length-test
-      return 0
-
-    clause2enum[clause] = enum
-    clauses.append((len(set(clause.args)), clause))
-
-  clauses = sorted(clauses, key=lambda x: x[0], reverse=True)
-  _, clauses = zip(*clauses)
-
-  for mapping in try_to_map([(c, clause2enum[c]) for c in clauses], {}):
-    if not mapping:
-      continue
-
-    checks_ok = True
-    for check in numerical_checks:
-      args = [mapping[a] for a in check.args]
-      if check.name == 'ncoll':
-        checks_ok = g.check_ncoll(args)
-      elif check.name == 'npara':
-        checks_ok = g.check_npara(args)
-      elif check.name == 'nperp':
-        checks_ok = g.check_nperp(args)
-      elif check.name == 'sameside':
-        checks_ok = g.check_sameside(args)
-      if not checks_ok:
-        break
-    if not checks_ok:
-      continue
-
-    yield mapping
-
-
-BUILT_IN_FNS = {
-    'cong_cong_cong_cyclic': match_cong_cong_cong_cyclic,
-    'cong_cong_cong_ncoll_contri*': match_cong_cong_cong_ncoll_contri,
-    'cong_cong_eqangle6_ncoll_contri*': match_cong_cong_eqangle6_ncoll_contri,
-    'eqangle6_eqangle6_ncoll_simtri': match_eqangle6_eqangle6_ncoll_simtri,
-    'eqangle6_eqangle6_ncoll_cong_contri': (
-        match_eqangle6_eqangle6_ncoll_cong_contri
-    ),  # pylint: disable=line-too-long
-    'eqangle6_eqangle6_ncoll_simtri2': match_eqangle6_eqangle6_ncoll_simtri2,
-    'eqangle6_eqangle6_ncoll_cong_contri2': (
-        match_eqangle6_eqangle6_ncoll_cong_contri2
-    ),  # pylint: disable=line-too-long
-    'eqratio6_eqratio6_ncoll_simtri*': match_eqratio6_eqratio6_ncoll_simtri,
-    'eqratio6_eqratio6_ncoll_cong_contri*': (
-        match_eqratio6_eqratio6_ncoll_cong_contri
-    ),  # pylint: disable=line-too-long
-    'eqangle_para': match_eqangle_para,
-    'eqangle_ncoll_cyclic': match_eqangle_ncoll_cyclic,
-    'eqratio6_eqangle6_ncoll_simtri*': match_eqratio6_eqangle6_ncoll_simtri,
-    'eqangle_perp_perp': match_eqangle_perp_perp,
-    'eqangle6_ncoll_cong': match_eqangle6_ncoll_cong,
-    'perp_perp_ncoll_para': match_perp_perp_ncoll_para,
-    'circle_perp_eqangle': match_circle_perp_eqangle,
-    'circle_eqangle_perp': match_circle_eqangle_perp,
-    'cyclic_eqangle_cong': match_cyclic_eqangle_cong,
-    'midp_perp_cong': match_midp_perp_cong,
-    'perp_perp_npara_eqangle': match_perp_perp_npara_eqangle,
-    'cyclic_eqangle': match_cyclic_eqangle,
-    'eqangle_eqangle_eqangle': match_eqangle_eqangle_eqangle,
-    'eqratio_eqratio_eqratio': match_eqratio_eqratio_eqratio,
-    'eqratio6_coll_ncoll_eqangle6': match_eqratio6_coll_ncoll_eqangle6,
-    'eqangle6_coll_ncoll_eqratio6': match_eqangle6_coll_ncoll_eqratio6,
-    'eqangle6_ncoll_cyclic': match_eqangle6_ncoll_cyclic,
-}
-
-
-SKIP_THEOREMS = set()
-
-
-def set_skip_theorems(theorems: set[str]) -> None:
-  SKIP_THEOREMS.update(theorems)
-
-
-MAX_BRANCH = 50_000
-
-
-def match_one_theorem(
-    g: gh.Graph,
-    cache: Callable[str, list[tuple[gm.Point, ...]]],
-    theorem: pr.Theorem
-) -> Generator[dict[str, gm.Point], None, None]:
-  """Match all instances of a single theorem (rule)."""
-  if cache is None:
-    cache = cache_match(g)
-
-  if theorem.name in SKIP_THEOREMS:
-    return []
-
-  if theorem.name.split('_')[-1] in SKIP_THEOREMS:
-    return []
-
-  if theorem.name in BUILT_IN_FNS:
-    mps = BUILT_IN_FNS[theorem.name](g, cache, theorem)
-  else:
-    mps = match_generic(g, cache, theorem)
-
-  mappings = []
-  for mp in mps:
-    mappings.append(mp)
-    if len(mappings) > MAX_BRANCH:  # cap branching at this number.
-      break
-
-  return mappings
-
-
-def match_all_theorems(
-    g: gh.Graph, theorems: list[pr.Theorem], goal: pr.Clause
-) -> dict[pr.Theorem, dict[pr.Theorem, dict[str, gm.Point]]]:
-  """Match all instances of all theorems (rules)."""
-  cache = cache_match(g)
-  # for BFS, collect all potential matches
-  # and then do it at the same time
-  theorem2mappings = {}
-
-  # Step 1: list all matches
-  for _, theorem in theorems.items():
-    name = theorem.name
-    if name.split('_')[-1] in [
-        'acompute',
-        'rcompute',
-        'fixl',
-        'fixc',
-        'fixb',
-        'fixt',
-        'fixp',
-    ]:
-      if goal and goal.name != name:
-        continue
-
-    mappings = match_one_theorem(g, cache, theorem)
-    if len(mappings):  # pylint: disable=g-explicit-length-test
-      theorem2mappings[theorem] = list(mappings)
-  return theorem2mappings
-
-
-def bfs_one_level(
-    g: gh.Graph,
-    theorems: list[pr.Theorem],
-    level: int,
-    controller: pr.Problem,
-    verbose: bool = False,
-    nm_check: bool = False,
-    timeout: int = 600,
-) -> tuple[
-    list[pr.Dependency],
-    dict[str, list[tuple[gm.Point, ...]]],
-    dict[str, list[tuple[gm.Point, ...]]],
-    int,
-]:
-  """Forward deduce one breadth-first level."""
-
-  # Step 1: match all theorems:
-  theorem2mappings = match_all_theorems(g, theorems, controller.goal)
-
-  # Step 2: traceback for each deduce:
-  theorem2deps = {}
-  t0 = time.time()
-  for theorem, mappings in theorem2mappings.items():
-    if time.time() - t0 > timeout:
-      break
-    mp_deps = []
-    for mp in mappings:
-      deps = EmptyDependency(level=level, rule_name=theorem.rule_name)
-      fail = False  # finding why deps might fail.
-
-      for p in theorem.premise:
-        p_args = [mp[a] for a in p.args]
-        # Trivial deps.
-        if p.name == 'cong':
-          a, b, c, d = p_args
-          if {a, b} == {c, d}:
-            continue
-        if p.name == 'para':
-          a, b, c, d = p_args
-          if {a, b} == {c, d}:
-            continue
-
-        if theorem.name in [
-            'cong_cong_eqangle6_ncoll_contri*',
-            'eqratio6_eqangle6_ncoll_simtri*',
-        ]:
-          if p.name in ['eqangle', 'eqangle6']:  # SAS or RAR
-            b, a, b, c, y, x, y, z = (  # pylint: disable=redeclared-assigned-name,unused-variable
-                p_args
-            )
-            if not nm.same_clock(a.num, b.num, c.num, x.num, y.num, z.num):
-              p_args = b, a, b, c, y, z, y, x
-
-        dep = Dependency(p.name, p_args, rule_name='', level=level)
-        try:
-          dep = dep.why_me_or_cache(g, level)
-        except:  # pylint: disable=bare-except
-          fail = True
-          break
-
-        if dep.why is None:
-          fail = True
-          break
-        g.cache_dep(p.name, p_args, dep)
-        deps.why.append(dep)
-
-      if fail:
-        continue
-
-      mp_deps.append((mp, deps))
-    theorem2deps[theorem] = mp_deps
-
-  theorem2deps = list(theorem2deps.items())
-
-  # Step 3: add conclusions to graph.
-  # Note that we do NOT mix step 2 and 3, strictly going for BFS.
-  added = []
-  for theorem, mp_deps in theorem2deps:
-    for mp, deps in mp_deps:
-      if time.time() - t0 > timeout:
-        break
-      name, args = theorem.conclusion_name_args(mp)
-      hash_conclusion = pr.hashed(name, args)
-      if hash_conclusion in g.cache:
-        continue
-
-      add = g.add_piece(name, args, deps=deps)
-      added += add
-
-  branching = len(added)
-
-  # Check if goal is found
-  if controller.goal:
-    args = []
-
-    for a in controller.goal.args:
-      if a in g._name2node:
-        a = g._name2node[a]
-      elif '/' in a:
-        a = create_consts_str(g, a)
-      elif a.isdigit():
-        a = int(a)
-      args.append(a)
-
-    if g.check(controller.goal.name, args):
-      return added, {}, {}, branching
-
-  # Run AR, but do NOT apply to the proof state (yet).
-  for dep in added:
-    g.add_algebra(dep, level)
-  derives, eq4s = g.derive_algebra(level, verbose=False)
-
-  branching += sum([len(x) for x in derives.values()])
-  branching += sum([len(x) for x in eq4s.values()])
-
-  return added, derives, eq4s, branching
-
-
-def create_consts_str(g: gh.Graph, s: str) -> gm.Angle | gm.Ratio:
-  if 'pi/' in s:
-    n, d = s.split('pi/')
-    n, d = int(n), int(d)
-    p0, _ = g.get_or_create_const_ang(n, d)
-  else:
-    n, d = s.split('/')
-    n, d = int(n), int(d)
-    p0, _ = g.get_or_create_const_rat(n, d)
-  return p0
-
-
-def do_algebra(
-    g: gh.Graph, added: list[pr.Dependency], verbose: bool = False
-) -> None:
-  for add in added:
-    g.add_algebra(add, None)
-  derives, eq4s = g.derive_algebra(level=None, verbose=verbose)
-  apply_derivations(g, derives)
-  apply_derivations(g, eq4s)
-
-
-def apply_derivations(
-    g: gh.Graph, derives: dict[str, list[tuple[gm.Point, ...]]]
-) -> list[pr.Dependency]:
-  applied = []
-  all_derives = list(derives.items())
-  for name, args in all_derives:
-    for arg in args:
-      applied += g.do_algebra(name, arg)
-  return applied
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Implements Deductive Database (DD)."""
+
+# pylint: disable=g-multiple-import,g-importing-member
+from collections import defaultdict
+import time
+from typing import Any, Callable, Generator
+
+import geometry as gm
+import graph as gh
+import graph_utils as utils
+import numericals as nm
+import problem as pr
+from problem import Dependency, EmptyDependency
+from typing import Union
+
+
+def intersect1(set1: set[Any], set2: set[Any]) -> Any:
+  for x in set1:
+    if x in set2:
+      return x
+  return None
+
+
+def diff_point(l: gm.Line, a: gm.Point) -> gm.Point:
+  for x in l.neighbors(gm.Point):
+    if x != a :
+      return x
+  return None
+
+
+# pylint: disable=protected-access
+# pylint: disable=unused-argument
+
+
+def match_eqratio_eqratio_eqratio(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqratio a b c d m n p q, eqratio c d e f p q r u => eqratio a b e f m n r u."""
+  for m1 in g.type2nodes[gm.Value]:
+    for m2 in g.type2nodes[gm.Value]:
+      rats1 = []
+      for rat in m1.neighbors(gm.Ratio):
+        l1, l2 = rat.lengths
+        if l1 is None or l2 is None:
+          continue
+        rats1.append((l1, l2))
+
+      rats2 = []
+      for rat in m2.neighbors(gm.Ratio):
+        l1, l2 = rat.lengths
+        if l1 is None or l2 is None:
+          continue
+        rats2.append((l1, l2))
+
+      pairs = []
+      for (l1, l2), (l3, l4) in utils.cross(rats1, rats2):
+        if l2 == l3:
+          pairs.append((l1, l2, l4))
+
+      for (l1, l12, l2), (l3, l34, l4) in utils.comb2(pairs):
+        if (l1, l12, l2) == (l3, l34, l4):
+          continue
+        if l1 == l2 or l3 == l4:
+          continue
+        if l1 == l12 or l12 == l2 or l3 == l34 or l4 == l34:
+          continue
+        # d12 - d1 = d34 - d3 = m1
+        # d2 - d12 = d4 - d34 = m2
+        # => d2 - d1 = d4 - d3 (= m1+m2)
+        a, b = g.two_points_of_length(l1)
+        c, d = g.two_points_of_length(l12)
+        m, n = g.two_points_of_length(l3)
+        p, q = g.two_points_of_length(l34)
+        # eqangle a b c d m n p q
+        e, f = g.two_points_of_length(l2)
+        r, u = g.two_points_of_length(l4)
+        yield dict(zip('abcdefmnpqru', [a, b, c, d, e, f, m, n, p, q, r, u]))
+
+
+def match_eqangle_eqangle_eqangle(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle a b c d m n p q, eqangle c d e f p q r u => eqangle a b e f m n r u."""
+  for m1 in g.type2nodes[gm.Measure]:
+    for m2 in g.type2nodes[gm.Measure]:
+      angs1 = []
+      for ang in m1.neighbors(gm.Angle):
+        d1, d2 = ang.directions
+        if d1 is None or d2 is None:
+          continue
+        angs1.append((d1, d2))
+
+      angs2 = []
+      for ang in m2.neighbors(gm.Angle):
+        d1, d2 = ang.directions
+        if d1 is None or d2 is None:
+          continue
+        angs2.append((d1, d2))
+
+      pairs = []
+      for (d1, d2), (d3, d4) in utils.cross(angs1, angs2):
+        if d2 == d3:
+          pairs.append((d1, d2, d4))
+
+      for (d1, d12, d2), (d3, d34, d4) in utils.comb2(pairs):
+        if (d1, d12, d2) == (d3, d34, d4):
+          continue
+        if d1 == d2 or d3 == d4:
+          continue
+        if d1 == d12 or d12 == d2 or d3 == d34 or d4 == d34:
+          continue
+        # d12 - d1 = d34 - d3 = m1
+        # d2 - d12 = d4 - d34 = m2
+        # => d2 - d1 = d4 - d3
+        a, b = g.two_points_on_direction(d1)
+        c, d = g.two_points_on_direction(d12)
+        m, n = g.two_points_on_direction(d3)
+        p, q = g.two_points_on_direction(d34)
+        # eqangle a b c d m n p q
+        e, f = g.two_points_on_direction(d2)
+        r, u = g.two_points_on_direction(d4)
+        yield dict(zip('abcdefmnpqru', [a, b, c, d, e, f, m, n, p, q, r, u]))
+
+
+def match_perp_perp_npara_eqangle(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match perp A B C D, perp E F G H, npara A B E F => eqangle A B E F C D G H."""
+  dpairs = []
+  for ang in g.vhalfpi.neighbors(gm.Angle):
+    d1, d2 = ang.directions
+    if d1 is None or d2 is None:
+      continue
+    dpairs.append((d1, d2))
+
+  for (d1, d2), (d3, d4) in utils.comb2(dpairs):
+    a, b = g.two_points_on_direction(d1)
+    c, d = g.two_points_on_direction(d2)
+    m, n = g.two_points_on_direction(d3)
+    p, q = g.two_points_on_direction(d4)
+    if g.check_npara([a, b, m, n]):
+      if ({a, b}, {c, d}) == ({m, n}, {p, q}):
+        continue
+      if ({a, b}, {c, d}) == ({p, q}, {m, n}):
+        continue
+
+      yield dict(zip('ABCDEFGH', [a, b, c, d, m, n, p, q]))
+
+
+def match_circle_coll_eqangle_midp(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match circle O A B C, coll M B C, eqangle A B A C O B O M => midp M B C."""
+  for p, a, b, c in g.all_circles():
+    ab = g._get_line(a, b)
+    if ab is None:
+      continue
+    if ab.val is None:
+      continue
+    ac = g._get_line(a, c)
+    if ac is None:
+      continue
+    if ac.val is None:
+      continue
+    pb = g._get_line(p, b)
+    if pb is None:
+      continue
+    if pb.val is None:
+      continue
+
+    bc = g._get_line(b, c)
+    if bc is None:
+      continue
+    bc_points = bc.neighbors(gm.Point, return_set=True)
+
+    anga, _ = g._get_angle(ab.val, ac.val)
+
+    for angp in pb.val.neighbors(gm.Angle):
+      if not g.is_equal(anga, angp):
+        continue
+
+      _, d = angp.directions
+      for l in d.neighbors(gm.Line):
+        l_points = l.neighbors(gm.Point, return_set=True)
+        m = intersect1(bc_points, l_points)
+        if m is not None:
+          yield dict(zip('ABCMO', [a, b, c, m, p]))
+
+
+def match_midp_perp_cong(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match midp M A B, perp O M A B => cong O A O B."""
+  for m, a, b in g.all_midps():
+    ab = g._get_line(a, b)
+    for l in m.neighbors(gm.Line):
+      if g.check_perpl(l, ab):
+        for o in l.neighbors(gm.Point):
+          if o != m:
+            yield dict(zip('ABMO', [a, b, m, o]))
+
+
+def match_cyclic_eqangle_cong(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match cyclic A B C P Q R, eqangle C A C B R P R Q => cong A B P Q."""
+  for c in g.type2nodes[gm.Circle]:
+    ps = c.neighbors(gm.Point)
+    for (a, b, c), (x, y, z) in utils.comb2(list(utils.perm3(ps))):
+      if {a, b, c} == {x, y, z}:
+        continue
+      if g.check_eqangle([c, a, c, b, z, x, z, y]):
+        yield dict(zip('ABCPQR', [a, b, c, x, y, z]))
+
+
+def match_circle_eqangle_perp(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match circle O A B C, eqangle A X A B C A C B => perp O A A X."""
+  for p, a, b, c in g.all_circles():
+    ca = g._get_line(c, a)
+    if ca is None:
+      continue
+    cb = g._get_line(c, b)
+    if cb is None:
+      continue
+    ab = g._get_line(a, b)
+    if ab is None:
+      continue
+
+    if ca.val is None:
+      continue
+    if cb.val is None:
+      continue
+    if ab.val is None:
+      continue
+
+    c_ang, _ = g._get_angle(cb.val, ca.val)
+    if c_ang is None:
+      continue
+
+    for ang in ab.val.neighbors(gm.Angle):
+      if g.is_equal(ang, c_ang):
+        _, d = ang.directions
+        for l in d.neighbors(gm.Line):
+          if a not in l.neighbors(gm.Point):
+            continue
+          x = diff_point(l, a)
+          if x is None:
+            continue
+          yield dict(zip('OABCX', [p, a, b, c, x]))
+        break
+
+
+def match_circle_perp_eqangle(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match circle O A B C, perp O A A X => eqangle A X A B C A C B."""
+  for p, a, b, c in g.all_circles():
+    pa = g._get_line(p, a)
+    if pa is None:
+      continue
+    if pa.val is None:
+      continue
+    for l in a.neighbors(gm.Line):
+      if g.check_perpl(pa, l):
+        x = diff_point(l, a)
+        if x is not None:
+          yield dict(zip('OABCX', [p, a, b, c, x]))
+
+def match_semicircle_eqangle_perp(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match circle O A B C, eqangle A X A B C A C B => perp O A A X."""
+  for p, a, b, c in g.all_circles():
+    ca = g._get_line(c, a)
+    if ca is None:
+      continue
+    cb = g._get_line(c, b)
+    if cb is None:
+      continue
+    ab = g._get_line(a, b)
+    if ab is None:
+      continue
+
+    if ca.val is None:
+      continue
+    if cb.val is None:
+      continue
+    if ab.val is None:
+      continue
+
+    c_ang, _ = g._get_angle(cb.val, ca.val)
+    if c_ang is None:
+      continue
+
+    for ang in ab.val.neighbors(gm.Angle):
+      if g.is_equal(ang, c_ang):
+        _, d = ang.directions
+        for l in d.neighbors(gm.Line):
+          if a not in l.neighbors(gm.Point):
+            continue
+          x = diff_point(l, a)
+          if x is None:
+            continue
+          yield dict(zip('OABCX', [p, a, b, c, x]))
+        break
+
+
+def match_semicircle_perp_eqangle(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+    """Match semicircle O A B, perp O A A X => eqangle A X A B O A B."""
+    for o, a, b, c in g.all_semicircles():
+      oa = g._get_line(o, a)
+      if oa is None:
+        continue
+      if oa.val is None:
+        continue
+      for l in a.neighbors(gm.Line):
+        if g.check_perpl(oa, l):
+          x = diff_point(l, a)
+          if x is not None:
+            yield dict(zip('OABCX', [o, a, b, c, x]))
+
+
+def match_perp_perp_ncoll_para(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match perp A B C D, perp C D E F, ncoll A B E => para A B E F."""
+  d2d = defaultdict(list)
+  for ang in g.vhalfpi.neighbors(gm.Angle):
+    d1, d2 = ang.directions
+    if d1 is None or d2 is None:
+      continue
+    d2d[d1] += [d2]
+    d2d[d2] += [d1]
+
+  for x, ys in d2d.items():
+    if len(ys) < 2:
+      continue
+    c, d = g.two_points_on_direction(x)
+    for y1, y2 in utils.comb2(ys):
+      a, b = g.two_points_on_direction(y1)
+      e, f = g.two_points_on_direction(y2)
+      if nm.check_ncoll([a.num, b.num, e.num]):
+        yield dict(zip('ABCDEF', [a, b, c, d, e, f]))
+
+
+def match_eqangle6_ncoll_cong(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 A O A B B A B O, ncoll O A B => cong O A O B."""
+  for a in g.type2nodes[gm.Point]:
+    for b, c in utils.comb2(g.type2nodes[gm.Point]):
+      if a == b or a == c:
+        continue
+      if g.check_eqangle([b, a, b, c, c, b, c, a]):
+        if g.check_ncoll([a, b, c]):
+          yield dict(zip('OAB', [a, b, c]))
+
+
+def match_eqangle_perp_perp(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle A B P Q C D U V, perp P Q U V => perp A B C D."""
+  for ang in g.vhalfpi.neighbors(gm.Angle):
+    # d1 perp d2
+    d1, d2 = ang.directions
+    if d1 is None or d2 is None:
+      continue
+    for d3, d4 in utils.comb2(g.type2nodes[gm.Direction]):
+      if d1 == d3 or d2 == d4:
+        continue
+      # if d1 - d3 = d2 - d4 => d3 perp d4
+      a13, a31 = g._get_angle(d1, d3)
+      a24, a42 = g._get_angle(d2, d4)
+      if a13 is None or a31 is None or a24 is None or a42 is None:
+        continue
+      if g.is_equal(a13, a24) and g.is_equal(a31, a42):
+        a, b = g.two_points_on_direction(d1)
+        c, d = g.two_points_on_direction(d2)
+        m, n = g.two_points_on_direction(d3)
+        p, q = g.two_points_on_direction(d4)
+        yield dict(zip('ABCDPQUV', [m, n, p, q, a, b, c, d]))
+
+
+def match_eqangle_ncoll_cyclic(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 P A P B Q A Q B, ncoll P Q A B => cyclic A B P Q."""
+  for l1, l2, l3, l4 in g.all_eqangles_distinct_linepairss():
+    if len(set([l1, l2, l3, l4])) < 4:
+      continue  # they all must be distinct.
+
+    p1s = l1.neighbors(gm.Point, return_set=True)
+    p2s = l2.neighbors(gm.Point, return_set=True)
+    p3s = l3.neighbors(gm.Point, return_set=True)
+    p4s = l4.neighbors(gm.Point, return_set=True)
+
+    p = intersect1(p1s, p2s)
+    if not p:
+      continue
+    q = intersect1(p3s, p4s)
+    if not q:
+      continue
+    a = intersect1(p1s, p3s)
+    if not a:
+      continue
+    b = intersect1(p2s, p4s)
+    if not b:
+      continue
+    if len(set([a, b, p, q])) < 4:
+      continue
+
+    if not g.check_ncoll([a, b, p, q]):
+      continue
+
+    yield dict(zip('ABPQ', [a, b, p, q]))
+
+
+def match_eqangle_para(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle A B P Q C D P Q => para A B C D."""
+  for measure in g.type2nodes[gm.Measure]:
+    angs = measure.neighbors(gm.Angle)
+    d12, d21 = defaultdict(list), defaultdict(list)
+    for ang in angs:
+      d1, d2 = ang.directions
+      if d1 is None or d2 is None:
+        continue
+      d12[d1].append(d2)
+      d21[d2].append(d1)
+
+    for d1, d2s in d12.items():
+      a, b = g.two_points_on_direction(d1)
+      for d2, d3 in utils.comb2(d2s):
+        c, d = g.two_points_on_direction(d2)
+        e, f = g.two_points_on_direction(d3)
+        yield dict(zip('ABCDPQ', [c, d, e, f, a, b]))
+
+
+def match_cyclic_eqangle(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match cyclic A B P Q => eqangle P A P B Q A Q B."""
+  record = set()
+  for a, b, c, d in g_matcher('cyclic'):
+    if (a, b, c, d) in record:
+      continue
+    record.add((a, b, c, d))
+    record.add((a, b, d, c))
+    record.add((b, a, c, d))
+    record.add((b, a, d, c))
+    yield dict(zip('ABPQ', [a, b, c, d]))
+
+
+def rotate_simtri(
+    a: gm.Point, b: gm.Point, c: gm.Point, x: gm.Point, y: gm.Point, z: gm.Point
+) -> Generator[tuple[gm.Point, ...], None, None]:
+  """Rotate points around for similar triangle predicates."""
+  yield (z, y, x, c, b, a)
+  for p in [
+      (b, c, a, y, z, x),
+      (c, a, b, z, x, y),
+      (x, y, z, a, b, c),
+      (y, z, x, b, c, a),
+      (z, x, y, c, a, b),
+  ]:
+    yield p
+    yield p[::-1]
+
+
+def match_cong_cong_cong_cyclic(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match cong O A O B, cong O B O C, cong O C O D => cyclic A B C D."""
+  for l in g.type2nodes[gm.Length]:
+    p2p = defaultdict(list)
+    for s in l.neighbors(gm.Segment):
+      a, b = s.points
+      p2p[a].append(b)
+      p2p[b].append(a)
+
+    for p, ps in p2p.items():
+      if len(ps) >= 4:
+        for a, b, c, d in utils.comb4(ps):
+          yield dict(zip('OABCD', [p, a, b, c, d]))
+
+
+def match_cong_cong_cong_ncoll_contri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match cong A B P Q, cong B C Q R, cong C A R P, ncoll A B C => contri* A B C P Q R."""
+  record = set()
+  for a, b, p, q in g_matcher('cong'):
+    for c in g.type2nodes[gm.Point]:
+      for r in g.type2nodes[gm.Point]:
+        if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
+          continue
+        if not g.check_ncoll([a, b, c]):
+          continue
+        if g.check_cong([b, c, q, r]) and g.check_cong([c, a, r, p]):
+          record.add((a, b, c, p, q, r))
+          yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
+
+
+def match_cong_cong_eqangle6_ncoll_contri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match cong A B P Q, cong B C Q R, eqangle6 B A B C Q P Q R, ncoll A B C => contri* A B C P Q R."""
+  record = set()
+  for a, b, p, q in g_matcher('cong'):
+    for c in g.type2nodes[gm.Point]:
+      if c in (a, b):
+        continue
+      for r in g.type2nodes[gm.Point]:
+        if r in (p, q):
+          continue
+
+        in_record = False
+        for x in [
+            (c, b, a, r, q, p),
+            (p, q, r, a, b, c),
+            (r, q, p, c, b, a),
+        ]:
+          if x in record:
+            in_record = True
+            break
+
+        if in_record:
+          continue
+
+        if not g.check_cong([b, c, q, r]):
+          continue
+        if not g.check_ncoll([a, b, c]):
+          continue
+
+        if nm.same_clock(a.num, b.num, c.num, p.num, q.num, r.num):
+          if g.check_eqangle([b, a, b, c, q, p, q, r]):
+            record.add((a, b, c, p, q, r))
+            yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
+        else:
+          if g.check_eqangle([b, a, b, c, q, r, q, p]):
+            record.add((a, b, c, p, q, r))
+            yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
+
+
+def match_eqratio6_eqangle6_ncoll_simtri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C => simtri* A B C P Q R."""
+  enums = g_matcher('eqratio6')
+
+  record = set()
+  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    if nm.same_clock(a.num, b.num, c.num, p.num, q.num, r.num):
+      if g.check_eqangle([b, a, b, c, q, p, q, r]):
+        record.add((a, b, c, p, q, r))
+        yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    elif g.check_eqangle([b, a, b, c, q, r, q, p]):
+      record.add((a, b, c, p, q, r))
+      yield dict(zip('ABCPQR', [a, b, c, p, q, r]))
+
+
+def match_eqangle6_eqangle6_ncoll_simtri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 B A B C Q P Q R, eqangle6 C A C B R P R Q, ncoll A B C => simtri A B C P Q R."""
+  enums = g_matcher('eqangle6')
+
+  record = set()
+  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqangle([c, a, c, b, r, p, r, q]):
+      continue
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def match_eqratio6_eqratio6_ncoll_simtri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C => simtri* A B C P Q R."""
+  enums = g_matcher('eqratio6')
+
+  record = set()
+  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqratio([c, a, c, b, r, p, r, q]):
+      continue
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def match_eqangle6_eqangle6_ncoll_simtri2(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 B A B C Q R Q P, eqangle6 C A C B R Q R P, ncoll A B C => simtri2 A B C P Q R."""
+  enums = g_matcher('eqangle6')
+
+  record = set()
+  for b, a, b, c, q, r, q, p in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_simtri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqangle([c, a, c, b, r, q, r, p]):
+      continue
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def rotate_contri(
+    a: gm.Point, b: gm.Point, c: gm.Point, x: gm.Point, y: gm.Point, z: gm.Point
+) -> Generator[tuple[gm.Point, ...], None, None]:
+  for p in [(b, a, c, y, x, z), (x, y, z, a, b, c), (y, x, z, b, a, c)]:
+    yield p
+
+
+def match_eqangle6_eqangle6_ncoll_cong_contri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 B A B C Q P Q R, eqangle6 C A C B R P R Q, ncoll A B C, cong A B P Q => contri A B C P Q R."""
+  enums = g_matcher('eqangle6')
+
+  record = set()
+  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if not g.check_cong([a, b, p, q]):
+      continue
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqangle([c, a, c, b, r, p, r, q]):
+      continue
+
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def match_eqratio6_eqratio6_ncoll_cong_contri(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C, cong A B P Q => contri* A B C P Q R."""
+  enums = g_matcher('eqratio6')
+
+  record = set()
+  for b, a, b, c, q, p, q, r in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if not g.check_cong([a, b, p, q]):
+      continue
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqratio([c, a, c, b, r, p, r, q]):
+      continue
+
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def match_eqangle6_eqangle6_ncoll_cong_contri2(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 B A B C Q R Q P, eqangle6 C A C B R Q R P, ncoll A B C, cong A B P Q => contri2 A B C P Q R."""
+  enums = g_matcher('eqangle6')
+
+  record = set()
+  for b, a, b, c, q, r, q, p in enums:  # pylint: disable=redeclared-assigned-name,unused-variable
+    if not g.check_cong([a, b, p, q]):
+      continue
+    if (a, b, c) == (p, q, r):
+      continue
+    if any([x in record for x in rotate_contri(a, b, c, p, q, r)]):
+      continue
+    if not g.check_eqangle([c, a, c, b, r, q, r, p]):
+      continue
+    if not g.check_ncoll([a, b, c]):
+      continue
+
+    mapping = dict(zip('ABCPQR', [a, b, c, p, q, r]))
+    record.add((a, b, c, p, q, r))
+    yield mapping
+
+
+def match_eqratio6_coll_ncoll_eqangle6(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqratio6 d b d c a b a c, coll d b c, ncoll a b c => eqangle6 a b a d a d a c."""
+  records = set()
+  for b, d, c in g_matcher('coll'):
+    for a in g.all_points():
+      if g.check_coll([a, b, c]):
+        continue
+      if (a, b, d, c) in records or (a, c, d, b) in records:
+        continue
+      records.add((a, b, d, c))
+
+      if g.check_eqratio([d, b, d, c, a, b, a, c]):
+        yield dict(zip('abcd', [a, b, c, d]))
+
+
+def match_eqangle6_coll_ncoll_eqratio6(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 a b a d a d a c, coll d b c, ncoll a b c => eqratio6 d b d c a b a c."""
+  records = set()
+  for b, d, c in g_matcher('coll'):
+    for a in g.all_points():
+      if g.check_coll([a, b, c]):
+        continue
+      if (a, b, d, c) in records or (a, c, d, b) in records:
+        continue
+      records.add((a, b, d, c))
+
+      if g.check_eqangle([a, b, a, d, a, d, a, c]):
+        yield dict(zip('abcd', [a, b, c, d]))
+
+
+def match_eqangle6_ncoll_cyclic(
+    g: gh.Graph,
+    g_matcher: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem,
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match eqangle6 P A P B Q A Q B, ncoll P Q A B => cyclic A B P Q."""
+  for a, b, a, c, x, y, x, z in g_matcher('eqangle6'):  # pylint: disable=redeclared-assigned-name,unused-variable
+    if (b, c) != (y, z) or a == x:
+      continue
+    if nm.check_ncoll([x.num for x in [a, b, c, x]]):
+      yield dict(zip('ABPQ', [b, c, a, x]))
+
+
+def match_all(
+    name: str, g: gh.Graph
+) -> Generator[tuple[gm.Point, ...], None, None]:
+  """Match all instances of a certain relation."""
+  if name in ['ncoll', 'npara', 'nperp']:
+    return []
+  if name == 'coll':
+    return g.all_colls()
+  if name == 'para':
+    return g.all_paras()
+  if name == 'perp':
+    return g.all_perps()
+  if name == 'cong':
+    return g.all_congs()
+  if name == 'eqangle':
+    return g.all_eqangles_8points()
+  if name == 'eqangle6':
+    return g.all_eqangles_6points()
+  if name == 'eqratio':
+    return g.all_eqratios_8points()
+  if name == 'eqratio6':
+    return g.all_eqratios_6points()
+  if name == 'cyclic':
+    return g.all_cyclics()
+  if name == 'midp':
+    return g.all_midps()
+  if name == 'circle':
+    return g.all_circles()
+  if name == 'semicircle':
+    return g.all_semicircles()
+  raise ValueError(f'Unrecognize {name}')
+
+
+def cache_match(
+    graph: gh.Graph,
+) -> Callable[str, list[tuple[gm.Point, ...]]]:
+  """Cache throughout one single BFS level."""
+  cache = {}
+
+  def match_fn(name: str) -> list[tuple[gm.Point, ...]]:
+    if name in cache:
+      return cache[name]
+
+    result = list(match_all(name, graph))
+    cache[name] = result
+    return result
+
+  return match_fn
+
+
+def try_to_map(
+    clause_enum: list[tuple[pr.Clause, list[tuple[gm.Point, ...]]]],
+    mapping: dict[str, gm.Point],
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Recursively try to match the remaining points given current mapping."""
+  if not clause_enum:
+    yield mapping
+    return
+
+  clause, enum = clause_enum[0]
+  for points in enum:
+    mpcpy = dict(mapping)
+
+    fail = False
+    for p, a in zip(points, clause.args):
+      if a in mpcpy and mpcpy[a] != p or p in mpcpy and mpcpy[p] != a:
+        fail = True
+        break
+      mpcpy[a] = p
+      mpcpy[p] = a
+
+    if fail:
+      continue
+
+    for m in try_to_map(clause_enum[1:], mpcpy):
+      yield m
+
+
+def match_generic(
+    g: gh.Graph,
+    cache: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match any generic rule that is not one of the above match_*() rules."""
+  clause2enum = {}
+
+  clauses = []
+  numerical_checks = []
+  for clause in theorem.premise:
+    if clause.name in ['ncoll', 'npara', 'nperp', 'sameside']:
+      numerical_checks.append(clause)
+      continue
+
+    enum = cache(clause.name)
+    if len(enum) == 0:  # pylint: disable=g-explicit-length-test
+      return 0
+
+    clause2enum[clause] = enum
+    clauses.append((len(set(clause.args)), clause))
+
+  clauses = sorted(clauses, key=lambda x: x[0], reverse=True)
+  _, clauses = zip(*clauses)
+
+  for mapping in try_to_map([(c, clause2enum[c]) for c in clauses], {}):
+    if not mapping:
+      continue
+
+    checks_ok = True
+    for check in numerical_checks:
+      args = [mapping[a] for a in check.args]
+      if check.name == 'ncoll':
+        checks_ok = g.check_ncoll(args)
+      elif check.name == 'npara':
+        checks_ok = g.check_npara(args)
+      elif check.name == 'nperp':
+        checks_ok = g.check_nperp(args)
+      elif check.name == 'sameside':
+        checks_ok = g.check_sameside(args)
+      if not checks_ok:
+        break
+    if not checks_ok:
+      continue
+
+    yield mapping
+
+
+BUILT_IN_FNS = {
+    'cong_cong_cong_cyclic': match_cong_cong_cong_cyclic,
+    'cong_cong_cong_ncoll_contri*': match_cong_cong_cong_ncoll_contri,
+    'cong_cong_eqangle6_ncoll_contri*': match_cong_cong_eqangle6_ncoll_contri,
+    'eqangle6_eqangle6_ncoll_simtri': match_eqangle6_eqangle6_ncoll_simtri,
+    'eqangle6_eqangle6_ncoll_cong_contri': (
+        match_eqangle6_eqangle6_ncoll_cong_contri
+    ),  # pylint: disable=line-too-long
+    'eqangle6_eqangle6_ncoll_simtri2': match_eqangle6_eqangle6_ncoll_simtri2,
+    'eqangle6_eqangle6_ncoll_cong_contri2': (
+        match_eqangle6_eqangle6_ncoll_cong_contri2
+    ),  # pylint: disable=line-too-long
+    'eqratio6_eqratio6_ncoll_simtri*': match_eqratio6_eqratio6_ncoll_simtri,
+    'eqratio6_eqratio6_ncoll_cong_contri*': (
+        match_eqratio6_eqratio6_ncoll_cong_contri
+    ),  # pylint: disable=line-too-long
+    'eqangle_para': match_eqangle_para,
+    'eqangle_ncoll_cyclic': match_eqangle_ncoll_cyclic,
+    'eqratio6_eqangle6_ncoll_simtri*': match_eqratio6_eqangle6_ncoll_simtri,
+    'eqangle_perp_perp': match_eqangle_perp_perp,
+    'eqangle6_ncoll_cong': match_eqangle6_ncoll_cong,
+    'perp_perp_ncoll_para': match_perp_perp_ncoll_para,
+    'circle_perp_eqangle': match_circle_perp_eqangle,
+    'circle_eqangle_perp': match_circle_eqangle_perp,
+    'cyclic_eqangle_cong': match_cyclic_eqangle_cong,
+    'midp_perp_cong': match_midp_perp_cong,
+    'perp_perp_npara_eqangle': match_perp_perp_npara_eqangle,
+    'cyclic_eqangle': match_cyclic_eqangle,
+    'eqangle_eqangle_eqangle': match_eqangle_eqangle_eqangle,
+    'eqratio_eqratio_eqratio': match_eqratio_eqratio_eqratio,
+    'eqratio6_coll_ncoll_eqangle6': match_eqratio6_coll_ncoll_eqangle6,
+    'eqangle6_coll_ncoll_eqratio6': match_eqangle6_coll_ncoll_eqratio6,
+    'eqangle6_ncoll_cyclic': match_eqangle6_ncoll_cyclic,
+    'semicircle_perp_eqangle': match_semicircle_perp_eqangle,
+    'semicircle_eqangle_perp': match_semicircle_eqangle_perp,
+}
+
+
+SKIP_THEOREMS = set()
+
+
+def set_skip_theorems(theorems: set[str]) -> None:
+  SKIP_THEOREMS.update(theorems)
+
+
+MAX_BRANCH = 50_000
+
+
+def match_one_theorem(
+    g: gh.Graph,
+    cache: Callable[str, list[tuple[gm.Point, ...]]],
+    theorem: pr.Theorem
+) -> Generator[dict[str, gm.Point], None, None]:
+  """Match all instances of a single theorem (rule)."""
+  if cache is None:
+    cache = cache_match(g)
+
+  if theorem.name in SKIP_THEOREMS:
+    return []
+
+  if theorem.name.split('_')[-1] in SKIP_THEOREMS:
+    return []
+
+  if theorem.name in BUILT_IN_FNS:
+    mps = BUILT_IN_FNS[theorem.name](g, cache, theorem)
+  else:
+    mps = match_generic(g, cache, theorem)
+
+  mappings = []
+  for mp in mps:
+    mappings.append(mp)
+    if len(mappings) > MAX_BRANCH:  # cap branching at this number.
+      break
+
+  return mappings
+
+
+def match_all_theorems(
+    g: gh.Graph, theorems: list[pr.Theorem], goal: pr.Clause
+) -> dict[pr.Theorem, dict[pr.Theorem, dict[str, gm.Point]]]:
+  """Match all instances of all theorems (rules)."""
+  cache = cache_match(g)
+  # for BFS, collect all potential matches
+  # and then do it at the same time
+  theorem2mappings = {}
+
+  # Step 1: list all matches
+  for _, theorem in theorems.items():
+    name = theorem.name
+    if name.split('_')[-1] in [
+        'acompute',
+        'rcompute',
+        'fixl',
+        'fixc',
+        'fixb',
+        'fixt',
+        'fixp',
+    ]:
+      if goal and goal.name != name:
+        continue
+
+    mappings = match_one_theorem(g, cache, theorem)
+    if len(mappings):  # pylint: disable=g-explicit-length-test
+      theorem2mappings[theorem] = list(mappings)
+  return theorem2mappings
+
+
+def bfs_one_level(
+    g: gh.Graph,
+    theorems: list[pr.Theorem],
+    level: int,
+    controller: pr.Problem,
+    verbose: bool = False,
+    nm_check: bool = False,
+    timeout: int = 600,
+) -> tuple[
+    list[pr.Dependency],
+    dict[str, list[tuple[gm.Point, ...]]],
+    dict[str, list[tuple[gm.Point, ...]]],
+    int,
+]:
+  """Forward deduce one breadth-first level."""
+
+  # Step 1: match all theorems:
+  theorem2mappings = match_all_theorems(g, theorems, controller.goal)
+
+  # Step 2: traceback for each deduce:
+  theorem2deps = {}
+  t0 = time.time()
+  for theorem, mappings in theorem2mappings.items():
+    if time.time() - t0 > timeout:
+      break
+    mp_deps = []
+    for mp in mappings:
+      deps = EmptyDependency(level=level, rule_name=theorem.rule_name)
+      fail = False  # finding why deps might fail.
+
+      for p in theorem.premise:
+        p_args = [mp[a] for a in p.args]
+        # Trivial deps.
+        if p.name == 'cong':
+          a, b, c, d = p_args
+          if {a, b} == {c, d}:
+            continue
+        if p.name == 'para':
+          a, b, c, d = p_args
+          if {a, b} == {c, d}:
+            continue
+
+        if theorem.name in [
+            'cong_cong_eqangle6_ncoll_contri*',
+            'eqratio6_eqangle6_ncoll_simtri*',
+        ]:
+          if p.name in ['eqangle', 'eqangle6']:  # SAS or RAR
+            b, a, b, c, y, x, y, z = (  # pylint: disable=redeclared-assigned-name,unused-variable
+                p_args
+            )
+            if not nm.same_clock(a.num, b.num, c.num, x.num, y.num, z.num):
+              p_args = b, a, b, c, y, z, y, x
+
+        dep = Dependency(p.name, p_args, rule_name='', level=level)
+        try:
+          dep = dep.why_me_or_cache(g, level)
+        except:  # pylint: disable=bare-except
+          fail = True
+          break
+
+        if dep.why is None:
+          fail = True
+          break
+        g.cache_dep(p.name, p_args, dep)
+        deps.why.append(dep)
+
+      if fail:
+        continue
+
+      mp_deps.append((mp, deps))
+    theorem2deps[theorem] = mp_deps
+
+  theorem2deps = list(theorem2deps.items())
+
+  # Step 3: add conclusions to graph.
+  # Note that we do NOT mix step 2 and 3, strictly going for BFS.
+  added = []
+  for theorem, mp_deps in theorem2deps:
+    for mp, deps in mp_deps:
+      if time.time() - t0 > timeout:
+        break
+      name, args = theorem.conclusion_name_args(mp)
+      hash_conclusion = pr.hashed(name, args)
+      if hash_conclusion in g.cache:
+        continue
+
+      add = g.add_piece(name, args, deps=deps)
+      added += add
+
+  branching = len(added)
+
+  # Check if goal is found
+  if controller.goal:
+    args = []
+
+    for a in controller.goal.args:
+      if a in g._name2node:
+        a = g._name2node[a]
+      elif '/' in a:
+        a = create_consts_str(g, a)
+      elif a.isdigit():
+        a = int(a)
+      args.append(a)
+
+    if g.check(controller.goal.name, args):
+      return added, {}, {}, branching
+
+  # Run AR, but do NOT apply to the proof state (yet).
+  for dep in added:
+    g.add_algebra(dep, level)
+  derives, eq4s = g.derive_algebra(level, verbose=False)
+
+  branching += sum([len(x) for x in derives.values()])
+  branching += sum([len(x) for x in eq4s.values()])
+
+  return added, derives, eq4s, branching
+
+
+def create_consts_str(g: gh.Graph, s: str) -> Union [gm.Angle, gm.Ratio]:
+  if 'pi/' in s:
+    n, d = s.split('pi/')
+    n, d = int(n), int(d)
+    p0, _ = g.get_or_create_const_ang(n, d)
+  else:
+    n, d = s.split('/')
+    n, d = int(n), int(d)
+    p0, _ = g.get_or_create_const_rat(n, d)
+  return p0
+
+
+def do_algebra(
+    g: gh.Graph, added: list[pr.Dependency], verbose: bool = False
+) -> None:
+  for add in added:
+    g.add_algebra(add, None)
+  derives, eq4s = g.derive_algebra(level=None, verbose=verbose)
+  apply_derivations(g, derives)
+  apply_derivations(g, eq4s)
+
+
+def apply_derivations(
+    g: gh.Graph, derives: dict[str, list[tuple[gm.Point, ...]]]
+) -> list[pr.Dependency]:
+  applied = []
+  all_derives = list(derives.items())
+  for name, args in all_derives:
+    for arg in args:
+      applied += g.do_algebra(name, arg)
+  return applied

ag4masses/alphageometry/ddar.py

@@ -1,159 +1,157 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Implements the combination DD+AR."""
-import time
-
-from absl import logging
-import dd
-import graph as gh
-import problem as pr
-from problem import Dependency  # pylint: disable=g-importing-member
-import trace_back
-
-
-def saturate_or_goal(
-    g: gh.Graph,
-    theorems: list[pr.Theorem],
-    level_times: list[float],
-    p: pr.Problem,
-    max_level: int = 100,
-    timeout: int = 600,
-) -> tuple[
-    list[dict[str, list[tuple[gh.Point, ...]]]],
-    list[dict[str, list[tuple[gh.Point, ...]]]],
-    list[int],
-    list[pr.Dependency],
-]:
-  """Run DD until saturation or goal found."""
-  derives = []
-  eq4s = []
-  branching = []
-  all_added = []
-
-  while len(level_times) < max_level:
-    level = len(level_times) + 1
-
-    t = time.time()
-    added, derv, eq4, n_branching = dd.bfs_one_level(
-        g, theorems, level, p, verbose=False, nm_check=True, timeout=timeout
-    )
-    all_added += added
-    branching.append(n_branching)
-
-    derives.append(derv)
-    eq4s.append(eq4)
-    level_time = time.time() - t
-
-    logging.info(f'Depth {level}/{max_level} time = {level_time}')  # pylint: disable=logging-fstring-interpolation
-    level_times.append(level_time)
-
-    if p.goal is not None:
-      goal_args = list(map(lambda x: g.get(x, lambda: int(x)), p.goal.args))
-      if g.check(p.goal.name, goal_args):  # found goal
-        break
-
-    if not added:  # saturated
-      break
-
-    if level_time > timeout:
-      break
-
-  return derives, eq4s, branching, all_added
-
-
-def solve(
-    g: gh.Graph,
-    theorems: list[pr.Problem],
-    controller: pr.Problem,
-    max_level: int = 1000,
-    timeout: int = 600,
-) -> tuple[gh.Graph, list[float], str, list[int], list[pr.Dependency]]:
-  """Alternate between DD and AR until goal is found."""
-  status = 'saturated'
-  level_times = []
-
-  dervs, eq4 = g.derive_algebra(level=0, verbose=False)
-  derives = [dervs]
-  eq4s = [eq4]
-  branches = []
-  all_added = []
-
-  while len(level_times) < max_level:
-    dervs, eq4, next_branches, added = saturate_or_goal(
-        g, theorems, level_times, controller, max_level, timeout=timeout
-    )
-    all_added += added
-
-    derives += dervs
-    eq4s += eq4
-    branches += next_branches
-
-    # Now, it is either goal or saturated
-    if controller.goal is not None:
-      goal_args = g.names2points(controller.goal.args)
-      if g.check(controller.goal.name, goal_args):  # found goal
-        status = 'solved'
-        break
-
-    if not derives:  # officially saturated.
-      logging.info("derives empty, breaking")
-      break
-
-    # Now we resort to algebra derivations.
-    added = []
-    while derives and not added:
-      added += dd.apply_derivations(g, derives.pop(0))
-
-    if added:
-      continue
-
-    # Final help from AR.
-    while eq4s and not added:
-      added += dd.apply_derivations(g, eq4s.pop(0))
-
-    all_added += added
-
-    if not added:  # Nothing left. saturated.
-      logging.info("Nothing added, breaking")
-      break
-
-  return g, level_times, status, branches, all_added
-
-
-def get_proof_steps(
-    g: gh.Graph, goal: pr.Clause, merge_trivials: bool = False
-) -> tuple[
-    list[pr.Dependency],
-    list[pr.Dependency],
-    list[tuple[list[pr.Dependency], list[pr.Dependency]]],
-    dict[tuple[str, ...], int],
-]:
-  """Extract proof steps from the built DAG."""
-  goal_args = g.names2nodes(goal.args)
-  query = Dependency(goal.name, goal_args, None, None)
-
-  setup, aux, log, setup_points = trace_back.get_logs(
-      query, g, merge_trivials=merge_trivials
-  )
-
-  refs = {}
-  setup = trace_back.point_log(setup, refs, set())
-  aux = trace_back.point_log(aux, refs, setup_points)
-
-  setup = [(prems, [tuple(p)]) for p, prems in setup]
-  aux = [(prems, [tuple(p)]) for p, prems in aux]
-
-  return setup, aux, log, refs
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Implements the combination DD+AR."""
+import time
+
+from absl import logging
+import dd
+import graph as gh
+import problem as pr
+from problem import Dependency  # pylint: disable=g-importing-member
+import trace_back
+
+
+def saturate_or_goal(
+    g: gh.Graph,
+    theorems: list[pr.Theorem],
+    level_times: list[float],
+    p: pr.Problem,
+    max_level: int = 100,
+    timeout: int = 600,
+) -> tuple[
+    list[dict[str, list[tuple[gh.Point, ...]]]],
+    list[dict[str, list[tuple[gh.Point, ...]]]],
+    list[int],
+    list[pr.Dependency],
+]:
+  """Run DD until saturation or goal found."""
+  derives = []
+  eq4s = []
+  branching = []
+  all_added = []
+
+  while len(level_times) < max_level:
+    level = len(level_times) + 1
+
+    t = time.time()
+    added, derv, eq4, n_branching = dd.bfs_one_level(
+        g, theorems, level, p, verbose=False, nm_check=True, timeout=timeout
+    )
+    all_added += added
+    branching.append(n_branching)
+
+    derives.append(derv)
+    eq4s.append(eq4)
+    level_time = time.time() - t
+
+    logging.info(f'Depth {level}/{max_level} time = {level_time}')  # pylint: disable=logging-fstring-interpolation
+    level_times.append(level_time)
+
+    if p.goal is not None:
+      goal_args = list(map(lambda x: g.get(x, lambda: int(x)), p.goal.args))
+      if g.check(p.goal.name, goal_args):  # found goal
+        break
+
+    if not added:  # saturated
+      break
+
+    if level_time > timeout:
+      break
+
+  return derives, eq4s, branching, all_added
+
+
+def solve(
+    g: gh.Graph,
+    theorems: list[pr.Problem],
+    controller: pr.Problem,
+    max_level: int = 1000,
+    timeout: int = 600,
+) -> tuple[gh.Graph, list[float], str, list[int], list[pr.Dependency]]:
+  """Alternate between DD and AR until goal is found."""
+  status = 'saturated'
+  level_times = []
+
+  dervs, eq4 = g.derive_algebra(level=0, verbose=False)
+  derives = [dervs]
+  eq4s = [eq4]
+  branches = []
+  all_added = []
+
+  while len(level_times) < max_level:
+    dervs, eq4, next_branches, added = saturate_or_goal(
+        g, theorems, level_times, controller, max_level, timeout=timeout
+    )
+    all_added += added
+
+    derives += dervs
+    eq4s += eq4
+    branches += next_branches
+
+    # Now, it is either goal or saturated
+    if controller.goal is not None:
+      goal_args = g.names2points(controller.goal.args)
+      if g.check(controller.goal.name, goal_args):  # found goal
+        status = 'solved'
+        break
+
+    if not derives:  # officially saturated.
+      break
+
+    # Now we resort to algebra derivations.
+    added = []
+    while derives and not added:
+      added += dd.apply_derivations(g, derives.pop(0))
+
+    if added:
+      continue
+
+    # Final help from AR.
+    while eq4s and not added:
+      added += dd.apply_derivations(g, eq4s.pop(0))
+
+    all_added += added
+
+    if not added:  # Nothing left. saturated.
+      break
+
+  return g, level_times, status, branches, all_added
+
+
+def get_proof_steps(
+    g: gh.Graph, goal: pr.Clause, merge_trivials: bool = False
+) -> tuple[
+    list[pr.Dependency],
+    list[pr.Dependency],
+    list[tuple[list[pr.Dependency], list[pr.Dependency]]],
+    dict[tuple[str, ...], int],
+]:
+  """Extract proof steps from the built DAG."""
+  goal_args = g.names2nodes(goal.args)
+  query = Dependency(goal.name, goal_args, None, None)
+
+  setup, aux, log, setup_points = trace_back.get_logs(
+      query, g, merge_trivials=merge_trivials
+  )
+
+  refs = {}
+  setup = trace_back.point_log(setup, refs, set())
+  aux = trace_back.point_log(aux, refs, setup_points)
+
+  setup = [(prems, [tuple(p)]) for p, prems in setup]
+  aux = [(prems, [tuple(p)]) for p, prems in aux]
+
+  return setup, aux, log, refs

ag4masses/alphageometry/decoder_stack.py

@@ -1,55 +1,55 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""The decoder stack in inference mode."""
-
-from typing import Any, Tuple
-
-import gin
-from transformer import decoder_stack
-import transformer_layer as tl
-
-
-struct = decoder_stack.struct
-nn_components = decoder_stack.nn_components
-position = decoder_stack.position
-jnp = decoder_stack.jnp
-attention = decoder_stack.attention
-
-DStackWindowState = decoder_stack.DStackWindowState
-
-Array = Any
-
-TransformerTaskConfig = decoder_stack.TransformerTaskConfig
-
-DStackDecoderState = Tuple[tl.DecoderState, ...]
-
-
-@gin.configurable
-class DecoderStackGenerate(decoder_stack.DecoderStack):
-  """Stack of transformer decoder layers."""
-
-  layer_factory = tl.TransformerLayerGenerate
-
-  def init_decoder_state_vanilla(
-      self, sequence_length: int, start_of_sequence: Array
-  ) -> DStackDecoderState:
-    """Return initial state for autoregressive generation."""
-    return tuple(
-        [
-            layer.init_decoder_state_vanilla(sequence_length, start_of_sequence)
-            for layer in self.transformer_layers
-        ]
-    )
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""The decoder stack in inference mode."""
+
+from typing import Any, Tuple
+
+import gin
+from meliad_lib.meliad.transformer import decoder_stack
+import transformer_layer as tl
+
+
+struct = decoder_stack.struct
+nn_components = decoder_stack.nn_components
+position = decoder_stack.position
+jnp = decoder_stack.jnp
+attention = decoder_stack.attention
+
+DStackWindowState = decoder_stack.DStackWindowState
+
+Array = Any
+
+TransformerTaskConfig = decoder_stack.TransformerTaskConfig
+
+DStackDecoderState = Tuple[tl.DecoderState, ...]
+
+
+@gin.configurable
+class DecoderStackGenerate(decoder_stack.DecoderStack):
+  """Stack of transformer decoder layers."""
+
+  layer_factory = tl.TransformerLayerGenerate
+
+  def init_decoder_state_vanilla(
+      self, sequence_length: int, start_of_sequence: Array
+  ) -> DStackDecoderState:
+    """Return initial state for autoregressive generation."""
+    return tuple(
+        [
+            layer.init_decoder_state_vanilla(sequence_length, start_of_sequence)
+            for layer in self.transformer_layers
+        ]
+    )

ag4masses/alphageometry/defs.txt

@@ -405,3 +405,15 @@ x : a b c
 a b c = ncoll a b c
 x : cyclic a b c x
 cyclic a b c
+
+semicircle x a b c
+x : a b c
+a b c = ncoll a b c
+x : cong x a x b; cong x b x c
+bline a b, bline a c
+
+on_semicircle x o a
+x : x o a
+o a = diff o a
+x : cong o x o a
+circle o o a

ag4masses/alphageometry/geometry.py

@@ -1,578 +1,621 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Implements geometric objects used in the graph representation."""
-from __future__ import annotations
-from collections import defaultdict  # pylint: disable=g-importing-member
-from typing import Any, Type
-
-# pylint: disable=protected-access
-
-
-class Node:
-  r"""Node in the proof state graph.
-
-  Can be Point, Line, Circle, etc.
-
-  Each node maintains a merge history to
-  other nodes if they are (found out to be) equivalent
-
-    a -> b -
-            \
-         c -> d -> e -> f -> g
-
-  d.merged_to = e
-  d.rep = g
-  d.merged_from = {a, b, c, d}
-  d.equivs = {a, b, c, d, e, f, g}
-  """
-
-  def __init__(self, name: str = '', graph: Any = None):
-    self.name = name or str(self)
-    self.graph = graph
-
-    self.edge_graph = {}
-    # Edge graph: what other nodes is connected to this node.
-    # edge graph = {
-    #   other1: {self1: deps, self2: deps},
-    #   other2: {self2: deps, self3: deps}
-    # }
-
-    self.merge_graph = {}
-    # Merge graph: history of merges with other nodes.
-    # merge_graph = {self1: {self2: deps1, self3: deps2}}
-
-    self.rep_by = None  # represented by.
-    self.members = {self}
-
-    self._val = None
-    self._obj = None
-
-    self.deps = []
-
-    # numerical representation.
-    self.num = None
-    self.change = set()  # what other nodes' num rely on this node?
-
-  def set_rep(self, node: Node) -> None:
-    if node == self:
-      return
-    self.rep_by = node
-    node.merge_edge_graph(self.edge_graph)
-    node.members.update(self.members)
-
-  def rep(self) -> Node:
-    x = self
-    while x.rep_by:
-      x = x.rep_by
-    return x
-
-  def why_rep(self) -> list[Any]:
-    return self.why_equal([self.rep()], None)
-
-  def rep_and_why(self) -> tuple[Node, list[Any]]:
-    rep = self.rep()
-    return rep, self.why_equal([rep], None)
-
-  def neighbors(
-      self, oftype: Type[Node], return_set: bool = False, do_rep: bool = True
-  ) -> list[Node]:
-    """Neighbors of this node in the proof state graph."""
-    if do_rep:
-      rep = self.rep()
-    else:
-      rep = self
-    result = set()
-
-    for n in rep.edge_graph:
-      if oftype is None or oftype and isinstance(n, oftype):
-        if do_rep:
-          result.add(n.rep())
-        else:
-          result.add(n)
-
-    if return_set:
-      return result
-    return list(result)
-
-  def merge_edge_graph(
-      self, new_edge_graph: dict[Node, dict[Node, list[Node]]]
-  ) -> None:
-    for x, xdict in new_edge_graph.items():
-      if x in self.edge_graph:
-        self.edge_graph[x].update(dict(xdict))
-      else:
-        self.edge_graph[x] = dict(xdict)
-
-  def merge(self, nodes: list[Node], deps: list[Any]) -> None:
-    for node in nodes:
-      self.merge_one(node, deps)
-
-  def merge_one(self, node: Node, deps: list[Any]) -> None:
-    node.rep().set_rep(self.rep())
-
-    if node in self.merge_graph:
-      return
-
-    self.merge_graph[node] = deps
-    node.merge_graph[self] = deps
-
-  def is_val(self, node: Node) -> bool:
-    return (
-        isinstance(self, Line)
-        and isinstance(node, Direction)
-        or isinstance(self, Segment)
-        and isinstance(node, Length)
-        or isinstance(self, Angle)
-        and isinstance(node, Measure)
-        or isinstance(self, Ratio)
-        and isinstance(node, Value)
-    )
-
-  def set_val(self, node: Node) -> None:
-    self._val = node
-
-  def set_obj(self, node: Node) -> None:
-    self._obj = node
-
-  @property
-  def val(self) -> Node:
-    if self._val is None:
-      return None
-    return self._val.rep()
-
-  @property
-  def obj(self) -> Node:
-    if self._obj is None:
-      return None
-    return self._obj.rep()
-
-  def equivs(self) -> set[Node]:
-    return self.rep().members
-
-  def connect_to(self, node: Node, deps: list[Any] = None) -> None:
-    rep = self.rep()
-
-    if node in rep.edge_graph:
-      rep.edge_graph[node].update({self: deps})
-    else:
-      rep.edge_graph[node] = {self: deps}
-
-    if self.is_val(node):
-      self.set_val(node)
-      node.set_obj(self)
-
-  def equivs_upto(self, level: int) -> dict[Node, Node]:
-    """What are the equivalent nodes up to a certain level."""
-    parent = {self: None}
-    visited = set()
-    queue = [self]
-    i = 0
-
-    while i < len(queue):
-      current = queue[i]
-      i += 1
-      visited.add(current)
-
-      for neighbor in current.merge_graph:
-        if (
-            level is not None
-            and current.merge_graph[neighbor].level is not None
-            and current.merge_graph[neighbor].level >= level
-        ):
-          continue
-        if neighbor not in visited:
-          queue.append(neighbor)
-          parent[neighbor] = current
-
-    return parent
-
-  def why_equal(self, others: list[Node], level: int) -> list[Any]:
-    """BFS why this node is equal to other nodes."""
-    others = set(others)
-    found = 0
-
-    parent = {}
-    queue = [self]
-    i = 0
-
-    while i < len(queue):
-      current = queue[i]
-      if current in others:
-        found += 1
-      if found == len(others):
-        break
-
-      i += 1
-
-      for neighbor in current.merge_graph:
-        if (
-            level is not None
-            and current.merge_graph[neighbor].level is not None
-            and current.merge_graph[neighbor].level >= level
-        ):
-          continue
-        if neighbor not in parent:
-          queue.append(neighbor)
-          parent[neighbor] = current
-
-    return bfs_backtrack(self, others, parent)
-
-  def why_equal_groups(
-      self, groups: list[list[Node]], level: int
-  ) -> tuple[list[Any], list[Node]]:
-    """BFS for why self is equal to at least one member of each group."""
-    others = [None for _ in groups]
-    found = 0
-
-    parent = {}
-    queue = [self]
-    i = 0
-
-    while i < len(queue):
-      current = queue[i]
-
-      for j, grp in enumerate(groups):
-        if others[j] is None and current in grp:
-          others[j] = current
-          found += 1
-
-      if found == len(others):
-        break
-
-      i += 1
-
-      for neighbor in current.merge_graph:
-        if (
-            level is not None
-            and current.merge_graph[neighbor].level is not None
-            and current.merge_graph[neighbor].level >= level
-        ):
-          continue
-        if neighbor not in parent:
-          queue.append(neighbor)
-          parent[neighbor] = current
-
-    return bfs_backtrack(self, others, parent), others
-
-  def why_val(self, level: int) -> list[Any]:
-    return self._val.why_equal([self.val], level)
-
-  def why_connect(self, node: Node, level: int = None) -> list[Any]:
-    rep = self.rep()
-    equivs = list(rep.edge_graph[node].keys())
-    if not equivs:
-      return None
-    equiv = equivs[0]
-    dep = rep.edge_graph[node][equiv]
-    return [dep] + self.why_equal(equiv, level)
-
-
-def why_connect(*pairs: list[tuple[Node, Node]]) -> list[Any]:
-  result = []
-  for node1, node2 in pairs:
-    result += node1.why_connect(node2)
-  return result
-
-
-def is_equiv(x: Node, y: Node, level: int = None) -> bool:
-  level = level or float('inf')
-  return x.why_equal([y], level) is not None
-
-
-def is_equal(x: Node, y: Node, level: int = None) -> bool:
-  if x == y:
-    return True
-  if x._val is None or y._val is None:
-    return False
-  if x.val != y.val:
-    return False
-  return is_equiv(x._val, y._val, level)
-
-
-def bfs_backtrack(
-    root: Node, leafs: list[Node], parent: dict[Node, Node]
-) -> list[Any]:
-  """Return the path given BFS trace of parent nodes."""
-  backtracked = {root}  # no need to backtrack further when touching this set.
-  deps = []
-  for node in leafs:
-    if node is None:
-      return None
-    if node in backtracked:
-      continue
-    if node not in parent:
-      return None
-    while node not in backtracked:
-      backtracked.add(node)
-      deps.append(node.merge_graph[parent[node]])
-      node = parent[node]
-
-  return deps
-
-
-class Point(Node):
-  pass
-
-
-class Line(Node):
-  """Node of type Line."""
-
-  def new_val(self) -> Direction:
-    return Direction()
-
-  def why_coll(self, points: list[Point], level: int = None) -> list[Any]:
-    """Why points are connected to self."""
-    level = level or float('inf')
-
-    groups = []
-    for p in points:
-      group = [
-          l
-          for l, d in self.edge_graph[p].items()
-          if d is None or d.level < level
-      ]
-      if not group:
-        return None
-      groups.append(group)
-
-    min_deps = None
-    for line in groups[0]:
-      deps, others = line.why_equal_groups(groups[1:], level)
-      if deps is None:
-        continue
-      for p, o in zip(points, [line] + others):
-        deps.append(self.edge_graph[p][o])
-      if min_deps is None or len(deps) < len(min_deps):
-        min_deps = deps
-
-    if min_deps is None:
-      return None
-    return [d for d in min_deps if d is not None]
-
-
-class Segment(Node):
-
-  def new_val(self) -> Length:
-    return Length()
-
-
-class Circle(Node):
-  """Node of type Circle."""
-
-  def why_cyclic(self, points: list[Point], level: int = None) -> list[Any]:
-    """Why points are connected to self."""
-    level = level or float('inf')
-
-    groups = []
-    for p in points:
-      group = [
-          c
-          for c, d in self.edge_graph[p].items()
-          if d is None or d.level < level
-      ]
-      if not group:
-        return None
-      groups.append(group)
-
-    min_deps = None
-    for circle in groups[0]:
-      deps, others = circle.why_equal_groups(groups[1:], level)
-      if deps is None:
-        continue
-      for p, o in zip(points, [circle] + others):
-        deps.append(self.edge_graph[p][o])
-
-      if min_deps is None or len(deps) < len(min_deps):
-        min_deps = deps
-
-    if min_deps is None:
-      return None
-    return [d for d in min_deps if d is not None]
-
-
-def why_equal(x: Node, y: Node, level: int = None) -> list[Any]:
-  if x == y:
-    return []
-  if not x._val or not y._val:
-    return None
-  if x._val == y._val:
-    return []
-  return x._val.why_equal([y._val], level)
-
-
-class Direction(Node):
-  pass
-
-
-def get_lines_thru_all(*points: list[Point]) -> list[Line]:
-  line2count = defaultdict(lambda: 0)
-  points = set(points)
-  for p in points:
-    for l in p.neighbors(Line):
-      line2count[l] += 1
-  return [l for l, count in line2count.items() if count == len(points)]
-
-
-def line_of_and_why(
-    points: list[Point], level: int = None
-) -> tuple[Line, list[Any]]:
-  """Why points are collinear."""
-  for l0 in get_lines_thru_all(*points):
-    for l in l0.equivs():
-      if all([p in l.edge_graph for p in points]):
-        x, y = l.points
-        colls = list({x, y} | set(points))
-        # if len(colls) < 3:
-        #   return l, []
-        why = l.why_coll(colls, level)
-        if why is not None:
-          return l, why
-
-  return None, None
-
-
-def get_circles_thru_all(*points: list[Point]) -> list[Circle]:
-  circle2count = defaultdict(lambda: 0)
-  points = set(points)
-  for p in points:
-    for c in p.neighbors(Circle):
-      circle2count[c] += 1
-  return [c for c, count in circle2count.items() if count == len(points)]
-
-
-def circle_of_and_why(
-    points: list[Point], level: int = None
-) -> tuple[Circle, list[Any]]:
-  """Why points are concyclic."""
-  for c0 in get_circles_thru_all(*points):
-    for c in c0.equivs():
-      if all([p in c.edge_graph for p in points]):
-        cycls = list(set(points))
-        why = c.why_cyclic(cycls, level)
-        if why is not None:
-          return c, why
-
-  return None, None
-
-
-def name_map(struct: Any) -> Any:
-  if isinstance(struct, list):
-    return [name_map(x) for x in struct]
-  elif isinstance(struct, tuple):
-    return tuple([name_map(x) for x in struct])
-  elif isinstance(struct, set):
-    return set([name_map(x) for x in struct])
-  elif isinstance(struct, dict):
-    return {name_map(x): name_map(y) for x, y in struct.items()}
-  else:
-    return getattr(struct, 'name', '')
-
-
-class Angle(Node):
-  """Node of type Angle."""
-
-  def new_val(self) -> Measure:
-    return Measure()
-
-  def set_directions(self, d1: Direction, d2: Direction) -> None:
-    self._d = d1, d2
-
-  @property
-  def directions(self) -> tuple[Direction, Direction]:
-    d1, d2 = self._d
-    if d1 is None or d2 is None:
-      return d1, d2
-    return d1.rep(), d2.rep()
-
-
-class Measure(Node):
-  pass
-
-
-class Length(Node):
-  pass
-
-
-class Ratio(Node):
-  """Node of type Ratio."""
-
-  def new_val(self) -> Value:
-    return Value()
-
-  def set_lengths(self, l1: Length, l2: Length) -> None:
-    self._l = l1, l2
-
-  @property
-  def lengths(self) -> tuple[Length, Length]:
-    l1, l2 = self._l
-    if l1 is None or l2 is None:
-      return l1, l2
-    return l1.rep(), l2.rep()
-
-
-class Value(Node):
-  pass
-
-
-def all_angles(
-    d1: Direction, d2: Direction, level: int = None
-) -> tuple[Angle, list[Direction], list[Direction]]:
-  level = level or float('inf')
-  d1s = d1.equivs_upto(level)
-  d2s = d2.equivs_upto(level)
-
-  for ang in d1.rep().neighbors(Angle):
-    d1_, d2_ = ang._d
-    if d1_ in d1s and d2_ in d2s:
-      yield ang, d1s, d2s
-
-
-def all_ratios(
-    d1, d2, level=None
-) -> tuple[Angle, list[Direction], list[Direction]]:
-  level = level or float('inf')
-  d1s = d1.equivs_upto(level)
-  d2s = d2.equivs_upto(level)
-
-  for ang in d1.rep().neighbors(Ratio):
-    d1_, d2_ = ang._l
-    if d1_ in d1s and d2_ in d2s:
-      yield ang, d1s, d2s
-
-
-RANKING = {
-    Point: 0,
-    Line: 1,
-    Segment: 2,
-    Circle: 3,
-    Direction: 4,
-    Length: 5,
-    Angle: 6,
-    Ratio: 7,
-    Measure: 8,
-    Value: 9,
-}
-
-
-def val_type(x: Node) -> Type[Node]:
-  if isinstance(x, Line):
-    return Direction
-  if isinstance(x, Segment):
-    return Length
-  if isinstance(x, Angle):
-    return Measure
-  if isinstance(x, Ratio):
-    return Value
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Implements geometric objects used in the graph representation."""
+from __future__ import annotations
+from collections import defaultdict  # pylint: disable=g-importing-member
+from typing import Any, Type
+import math
+# pylint: disable=protected-access
+
+
+class Node:
+  r"""Node in the proof state graph.
+
+  Can be Point, Line, Circle, etc.
+
+  Each node maintains a merge history to
+  other nodes if they are (found out to be) equivalent
+
+    a -> b -
+            \
+         c -> d -> e -> f -> g
+
+  d.merged_to = e
+  d.rep = g
+  d.merged_from = {a, b, c, d}
+  d.equivs = {a, b, c, d, e, f, g}
+  """
+
+  def __init__(self, name: str = '', graph: Any = None):
+    self.name = name or str(self)
+    self.graph = graph
+
+    self.edge_graph = {}
+    # Edge graph: what other nodes is connected to this node.
+    # edge graph = {
+    #   other1: {self1: deps, self2: deps},
+    #   other2: {self2: deps, self3: deps}
+    # }
+
+    self.merge_graph = {}
+    # Merge graph: history of merges with other nodes.
+    # merge_graph = {self1: {self2: deps1, self3: deps2}}
+
+    self.rep_by = None  # represented by.
+    self.members = {self}
+
+    self._val = None
+    self._obj = None
+
+    self.deps = []
+
+    # numerical representation.
+    self.num = None
+    self.change = set()  # what other nodes' num rely on this node?
+
+  def set_rep(self, node: Node) -> None:
+    if node == self:
+      return
+    self.rep_by = node
+    node.merge_edge_graph(self.edge_graph)
+    node.members.update(self.members)
+
+  def rep(self) -> Node:
+    x = self
+    while x.rep_by:
+      x = x.rep_by
+    return x
+
+  def why_rep(self) -> list[Any]:
+    return self.why_equal([self.rep()], None)
+
+  def rep_and_why(self) -> tuple[Node, list[Any]]:
+    rep = self.rep()
+    return rep, self.why_equal([rep], None)
+
+  def neighbors(
+      self, oftype: Type[Node], return_set: bool = False, do_rep: bool = True
+  ) -> list[Node]:
+    """Neighbors of this node in the proof state graph."""
+    if do_rep:
+      rep = self.rep()
+    else:
+      rep = self
+    result = set()
+
+    for n in rep.edge_graph:
+      if oftype is None or oftype and isinstance(n, oftype):
+        if do_rep:
+          result.add(n.rep())
+        else:
+          result.add(n)
+
+    if return_set:
+      return result
+    return list(result)
+
+  def merge_edge_graph(
+      self, new_edge_graph: dict[Node, dict[Node, list[Node]]]
+  ) -> None:
+    for x, xdict in new_edge_graph.items():
+      if x in self.edge_graph:
+        self.edge_graph[x].update(dict(xdict))
+      else:
+        self.edge_graph[x] = dict(xdict)
+
+  def merge(self, nodes: list[Node], deps: list[Any]) -> None:
+    for node in nodes:
+      self.merge_one(node, deps)
+
+  def merge_one(self, node: Node, deps: list[Any]) -> None:
+    node.rep().set_rep(self.rep())
+
+    if node in self.merge_graph:
+      return
+
+    self.merge_graph[node] = deps
+    node.merge_graph[self] = deps
+
+  def is_val(self, node: Node) -> bool:
+    return (
+        isinstance(self, Line)
+        and isinstance(node, Direction)
+        or isinstance(self, Segment)
+        and isinstance(node, Length)
+        or isinstance(self, Angle)
+        and isinstance(node, Measure)
+        or isinstance(self, Ratio)
+        and isinstance(node, Value)
+    )
+
+  def set_val(self, node: Node) -> None:
+    self._val = node
+
+  def set_obj(self, node: Node) -> None:
+    self._obj = node
+
+  @property
+  def val(self) -> Node:
+    if self._val is None:
+      return None
+    return self._val.rep()
+
+  @property
+  def obj(self) -> Node:
+    if self._obj is None:
+      return None
+    return self._obj.rep()
+
+  def equivs(self) -> set[Node]:
+    return self.rep().members
+
+  def connect_to(self, node: Node, deps: list[Any] = None) -> None:
+    rep = self.rep()
+
+    if node in rep.edge_graph:
+      rep.edge_graph[node].update({self: deps})
+    else:
+      rep.edge_graph[node] = {self: deps}
+
+    if self.is_val(node):
+      self.set_val(node)
+      node.set_obj(self)
+
+  def equivs_upto(self, level: int) -> dict[Node, Node]:
+    """What are the equivalent nodes up to a certain level."""
+    parent = {self: None}
+    visited = set()
+    queue = [self]
+    i = 0
+
+    while i < len(queue):
+      current = queue[i]
+      i += 1
+      visited.add(current)
+
+      for neighbor in current.merge_graph:
+        if (
+            level is not None
+            and current.merge_graph[neighbor].level is not None
+            and current.merge_graph[neighbor].level >= level
+        ):
+          continue
+        if neighbor not in visited:
+          queue.append(neighbor)
+          parent[neighbor] = current
+
+    return parent
+
+  def why_equal(self, others: list[Node], level: int) -> list[Any]:
+    """BFS why this node is equal to other nodes."""
+    others = set(others)
+    found = 0
+
+    parent = {}
+    queue = [self]
+    i = 0
+
+    while i < len(queue):
+      current = queue[i]
+      if current in others:
+        found += 1
+      if found == len(others):
+        break
+
+      i += 1
+
+      for neighbor in current.merge_graph:
+        if (
+            level is not None
+            and current.merge_graph[neighbor].level is not None
+            and current.merge_graph[neighbor].level >= level
+        ):
+          continue
+        if neighbor not in parent:
+          queue.append(neighbor)
+          parent[neighbor] = current
+
+    return bfs_backtrack(self, others, parent)
+
+  def why_equal_groups(
+      self, groups: list[list[Node]], level: int
+  ) -> tuple[list[Any], list[Node]]:
+    """BFS for why self is equal to at least one member of each group."""
+    others = [None for _ in groups]
+    found = 0
+
+    parent = {}
+    queue = [self]
+    i = 0
+
+    while i < len(queue):
+      current = queue[i]
+
+      for j, grp in enumerate(groups):
+        if others[j] is None and current in grp:
+          others[j] = current
+          found += 1
+
+      if found == len(others):
+        break
+
+      i += 1
+
+      for neighbor in current.merge_graph:
+        if (
+            level is not None
+            and current.merge_graph[neighbor].level is not None
+            and current.merge_graph[neighbor].level >= level
+        ):
+          continue
+        if neighbor not in parent:
+          queue.append(neighbor)
+          parent[neighbor] = current
+
+    return bfs_backtrack(self, others, parent), others
+
+  def why_val(self, level: int) -> list[Any]:
+    return self._val.why_equal([self.val], level)
+
+  def why_connect(self, node: Node, level: int = None) -> list[Any]:
+    rep = self.rep()
+    equivs = list(rep.edge_graph[node].keys())
+    if not equivs:
+      return None
+    equiv = equivs[0]
+    dep = rep.edge_graph[node][equiv]
+    return [dep] + self.why_equal(equiv, level)
+
+
+def why_connect(*pairs: list[tuple[Node, Node]]) -> list[Any]:
+  result = []
+  for node1, node2 in pairs:
+    result += node1.why_connect(node2)
+  return result
+
+
+def is_equiv(x: Node, y: Node, level: int = None) -> bool:
+  level = level or float('inf')
+  return x.why_equal([y], level) is not None
+
+
+def is_equal(x: Node, y: Node, level: int = None) -> bool:
+  if x == y:
+    return True
+  if x._val is None or y._val is None:
+    return False
+  if x.val != y.val:
+    return False
+  return is_equiv(x._val, y._val, level)
+
+
+def bfs_backtrack(
+    root: Node, leafs: list[Node], parent: dict[Node, Node]
+) -> list[Any]:
+  """Return the path given BFS trace of parent nodes."""
+  backtracked = {root}  # no need to backtrack further when touching this set.
+  deps = []
+  for node in leafs:
+    if node is None:
+      return None
+    if node in backtracked:
+      continue
+    if node not in parent:
+      return None
+    while node not in backtracked:
+      backtracked.add(node)
+      deps.append(node.merge_graph[parent[node]])
+      node = parent[node]
+
+  return deps
+
+
+
+class Point(Node):
+  pass
+
+
+class Line(Node):
+  """Node of type Line."""
+
+  def new_val(self) -> Direction:
+    return Direction()
+
+  def why_coll(self, points: list[Point], level: int = None) -> list[Any]:
+    """Why points are connected to self."""
+    level = level or float('inf')
+
+    groups = []
+    for p in points:
+      group = [
+          l
+          for l, d in self.edge_graph[p].items()
+          if d is None or d.level < level
+      ]
+      if not group:
+        return None
+      groups.append(group)
+
+    min_deps = None
+    for line in groups[0]:
+      deps, others = line.why_equal_groups(groups[1:], level)
+      if deps is None:
+        continue
+      for p, o in zip(points, [line] + others):
+        deps.append(self.edge_graph[p][o])
+      if min_deps is None or len(deps) < len(min_deps):
+        min_deps = deps
+
+    if min_deps is None:
+      return None
+    return [d for d in min_deps if d is not None]
+
+
+class Segment(Node):
+
+  def new_val(self) -> Length:
+    return Length()
+
+
+class Circle(Node): 
+  """Node of type Circle."""
+
+  def why_cyclic(self, points: list[Point], level: int = None) -> list[Any]:
+    """Why points are connected to self."""
+    level = level or float('inf')
+
+    groups = []
+    for p in points:
+      group = [
+          c
+          for c, d in self.edge_graph[p].items()
+          if d is None or d.level < level
+      ]
+      if not group:
+        return None
+      groups.append(group)
+
+    min_deps = None
+    for circle in groups[0]:
+      deps, others = circle.why_equal_groups(groups[1:], level)
+      if deps is None:
+        continue
+      for p, o in zip(points, [circle] + others):
+        deps.append(self.edge_graph[p][o])
+
+      if min_deps is None or len(deps) < len(min_deps):
+        min_deps = deps
+
+    if min_deps is None:
+      return None
+    return [d for d in min_deps if d is not None]
+
+# geometry.py
+class SemiCircle(Circle):
+    """Node of type SemiCircle, inheriting from Circle."""
+
+    def __init__(self, center: Point, radius: float):
+        """Initialize a semicircle with a center and radius."""
+        super().__init__(center, radius)
+
+    def contains_point(self, point: Point) -> bool:
+        """Check if a point lies inside the semicircle."""
+        # Check if point lies within the radius distance from the center (circle constraint)
+        if point.distance(self.center) > self.radius:
+            return False
+
+        # Additional logic to determine if the point is within the semicircle
+        return self.is_on_correct_side(point)
+
+    def is_on_correct_side(self, point: Point) -> bool:
+        """Check if the point is on the correct side of the semicircle."""
+        # Calculate the angle between the center and the point
+        angle = math.atan2(point.y - self.center.y, point.x - self.center.x)
+
+        # Determine the boundary angles of the semicircle
+        # Assuming the semicircle is oriented horizontally with the flat side down
+        start_angle = -math.pi / 2
+        end_angle = math.pi / 2
+
+        # Check if the point's angle lies within the boundary angles
+        return start_angle <= angle <= end_angle
+    
+    def why_cyclic(self, points: list[Point], level: int = None) -> list[Any]:
+        """Override why_cyclic to apply semicircle constraints."""
+        cyclic_points = super().why_cyclic(points, level)
+        if cyclic_points is None:
+            return None
+
+        # Ensure that all points lie within the semicircle
+        if all(self.contains_point(p) for p in points):
+            return cyclic_points
+        return None
+
+def why_equal(x: Node, y: Node, level: int = None) -> list[Any]:
+  if x == y:
+    return []
+  if not x._val or not y._val:
+    return None
+  if x._val == y._val:
+    return []
+  return x._val.why_equal([y._val], level)
+
+
+
+class Direction(Node):
+  pass
+
+
+def get_lines_thru_all(*points: list[Point]) -> list[Line]:
+  line2count = defaultdict(lambda: 0)
+  points = set(points)
+  for p in points:
+    for l in p.neighbors(Line):
+      line2count[l] += 1
+  return [l for l, count in line2count.items() if count == len(points)]
+
+
+def line_of_and_why(
+    points: list[Point], level: int = None
+) -> tuple[Line, list[Any]]:
+  """Why points are collinear."""
+  for l0 in get_lines_thru_all(*points):
+    for l in l0.equivs():
+      if all([p in l.edge_graph for p in points]):
+        x, y = l.points
+        colls = list({x, y} | set(points))
+        # if len(colls) < 3:
+        #   return l, []
+        why = l.why_coll(colls, level)
+        if why is not None:
+          return l, why
+
+  return None, None
+
+
+def get_circles_thru_all(*points: list[Point]) -> list[Circle]:
+  circle2count = defaultdict(lambda: 0)
+  points = set(points)
+  for p in points:
+    for c in p.neighbors(Circle):
+      circle2count[c] += 1
+  return [c for c, count in circle2count.items() if count == len(points)]
+
+
+def circle_of_and_why(
+    points: list[Point], level: int = None
+) -> tuple[Circle, list[Any]]:
+  """Why points are concyclic."""
+  for c0 in get_circles_thru_all(*points):
+    for c in c0.equivs():
+      if all([p in c.edge_graph for p in points]):
+        cycls = list(set(points))
+        why = c.why_cyclic(cycls, level)
+        if why is not None:
+          return c, why
+
+  return None, None
+
+
+def name_map(struct: Any) -> Any:
+  if isinstance(struct, list):
+    return [name_map(x) for x in struct]
+  elif isinstance(struct, tuple):
+    return tuple([name_map(x) for x in struct])
+  elif isinstance(struct, set):
+    return set([name_map(x) for x in struct])
+  elif isinstance(struct, dict):
+    return {name_map(x): name_map(y) for x, y in struct.items()}
+  else:
+    return getattr(struct, 'name', '')
+
+
+class Angle(Node):
+  """Node of type Angle."""
+
+  def new_val(self) -> Measure:
+    return Measure()
+
+  def set_directions(self, d1: Direction, d2: Direction) -> None:
+    self._d = d1, d2
+
+  @property
+  def directions(self) -> tuple[Direction, Direction]:
+    d1, d2 = self._d
+    if d1 is None or d2 is None:
+      return d1, d2
+    return d1.rep(), d2.rep()
+
+
+class Measure(Node):
+  pass
+
+
+class Length(Node):
+  pass
+
+
+class Ratio(Node):
+  """Node of type Ratio."""
+
+  def new_val(self) -> Value:
+    return Value()
+
+  def set_lengths(self, l1: Length, l2: Length) -> None:
+    self._l = l1, l2
+
+  @property
+  def lengths(self) -> tuple[Length, Length]:
+    l1, l2 = self._l
+    if l1 is None or l2 is None:
+      return l1, l2
+    return l1.rep(), l2.rep()
+
+
+class Value(Node):
+  pass
+
+
+def all_angles(
+    d1: Direction, d2: Direction, level: int = None
+) -> tuple[Angle, list[Direction], list[Direction]]:
+  level = level or float('inf')
+  d1s = d1.equivs_upto(level)
+  d2s = d2.equivs_upto(level)
+
+  for ang in d1.rep().neighbors(Angle):
+    d1_, d2_ = ang._d
+    if d1_ in d1s and d2_ in d2s:
+      yield ang, d1s, d2s
+
+
+def all_ratios(
+    d1, d2, level=None
+) -> tuple[Angle, list[Direction], list[Direction]]:
+  level = level or float('inf')
+  d1s = d1.equivs_upto(level)
+  d2s = d2.equivs_upto(level)
+
+  for ang in d1.rep().neighbors(Ratio):
+    d1_, d2_ = ang._l
+    if d1_ in d1s and d2_ in d2s:
+      yield ang, d1s, d2s
+
+
+RANKING = {
+    Point: 0,
+    Line: 1,
+    Segment: 2,
+    Circle: 3,
+    SemiCircle: 3,
+    Direction: 4,
+    Length: 5,
+    Angle: 6,
+    Ratio: 7,
+    Measure: 8,
+    Value: 9,
+}
+
+
+def val_type(x: Node) -> Type[Node]:
+  if isinstance(x, Line):
+    return Direction
+  if isinstance(x, Segment):
+    return Length
+  if isinstance(x, Angle):
+    return Measure
+  if isinstance(x, Ratio):
+    return Value

ag4masses/alphageometry/graph_utils.py

@@ -1,132 +1,132 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Utilizations for graph representation.
-
-Mainly for listing combinations and permutations of elements.
-"""
-
-from geometry import Point
-
-
-def _cross(elems1, elems2):
-  for e1 in elems1:
-    for e2 in elems2:
-      yield e1, e2
-
-
-def cross(elems1, elems2):
-  return list(_cross(elems1, elems2))
-
-
-def _comb2(elems):
-  if len(elems) < 2:
-    return
-  for i, e1 in enumerate(elems[:-1]):
-    for e2 in elems[i + 1 :]:
-      yield e1, e2
-
-
-def comb2(elems):
-  return list(_comb2(elems))
-
-
-def _comb3(elems):
-  if len(elems) < 3:
-    return
-  for i, e1 in enumerate(elems[:-2]):
-    for j, e2 in enumerate(elems[i + 1 : -1]):
-      for e3 in elems[i + j + 2 :]:
-        yield e1, e2, e3
-
-
-def comb3(elems):
-  return list(_comb3(elems))
-
-
-def _comb4(elems):
-  if len(elems) < 4:
-    return
-  for i, e1 in enumerate(elems[:-3]):
-    for j, e2 in enumerate(elems[i + 1 : -2]):
-      for e3, e4 in _comb2(elems[i + j + 2 :]):
-        yield e1, e2, e3, e4
-
-
-def comb4(elems):
-  return list(_comb4(elems))
-
-
-def _perm2(elems):
-  for e1, e2 in comb2(elems):
-    yield e1, e2
-    yield e2, e1
-
-
-def perm2(elems):
-  return list(_perm2(elems))
-
-
-def _all_4points(l1, l2):
-  p1s = l1.neighbors(Point)
-  p2s = l2.neighbors(Point)
-  for a, b in perm2(p1s):
-    for c, d in perm2(p2s):
-      yield a, b, c, d
-
-
-def all_4points(l1, l2):
-  return list(_all_4points(l1, l2))
-
-
-def _all_8points(l1, l2, l3, l4):
-  for a, b, c, d in all_4points(l1, l2):
-    for e, f, g, h in all_4points(l3, l4):
-      yield (a, b, c, d, e, f, g, h)
-
-
-def all_8points(l1, l2, l3, l4):
-  return list(_all_8points(l1, l2, l3, l4))
-
-
-def _perm3(elems):
-  for x in elems:
-    for y in elems:
-      if y == x:
-        continue
-      for z in elems:
-        if z not in (x, y):
-          yield x, y, z
-
-
-def perm3(elems):
-  return list(_perm3(elems))
-
-
-def _perm4(elems):
-  for x in elems:
-    for y in elems:
-      if y == x:
-        continue
-      for z in elems:
-        if z in (x, y):
-          continue
-        for t in elems:
-          if t not in (x, y, z):
-            yield x, y, z, t
-
-
-def perm4(elems):
-  return list(_perm4(elems))
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utilizations for graph representation.
+
+Mainly for listing combinations and permutations of elements.
+"""
+
+from geometry import Point
+
+
+def _cross(elems1, elems2):
+  for e1 in elems1:
+    for e2 in elems2:
+      yield e1, e2
+
+
+def cross(elems1, elems2):
+  return list(_cross(elems1, elems2))
+
+
+def _comb2(elems):
+  if len(elems) < 2:
+    return
+  for i, e1 in enumerate(elems[:-1]):
+    for e2 in elems[i + 1 :]:
+      yield e1, e2
+
+
+def comb2(elems):
+  return list(_comb2(elems))
+
+
+def _comb3(elems):
+  if len(elems) < 3:
+    return
+  for i, e1 in enumerate(elems[:-2]):
+    for j, e2 in enumerate(elems[i + 1 : -1]):
+      for e3 in elems[i + j + 2 :]:
+        yield e1, e2, e3
+
+
+def comb3(elems):
+  return list(_comb3(elems))
+
+
+def _comb4(elems):
+  if len(elems) < 4:
+    return
+  for i, e1 in enumerate(elems[:-3]):
+    for j, e2 in enumerate(elems[i + 1 : -2]):
+      for e3, e4 in _comb2(elems[i + j + 2 :]):
+        yield e1, e2, e3, e4
+
+
+def comb4(elems):
+  return list(_comb4(elems))
+
+
+def _perm2(elems):
+  for e1, e2 in comb2(elems):
+    yield e1, e2
+    yield e2, e1
+
+
+def perm2(elems):
+  return list(_perm2(elems))
+
+
+def _all_4points(l1, l2):
+  p1s = l1.neighbors(Point)
+  p2s = l2.neighbors(Point)
+  for a, b in perm2(p1s):
+    for c, d in perm2(p2s):
+      yield a, b, c, d
+
+
+def all_4points(l1, l2):
+  return list(_all_4points(l1, l2))
+
+
+def _all_8points(l1, l2, l3, l4):
+  for a, b, c, d in all_4points(l1, l2):
+    for e, f, g, h in all_4points(l3, l4):
+      yield (a, b, c, d, e, f, g, h)
+
+
+def all_8points(l1, l2, l3, l4):
+  return list(_all_8points(l1, l2, l3, l4))
+
+
+def _perm3(elems):
+  for x in elems:
+    for y in elems:
+      if y == x:
+        continue
+      for z in elems:
+        if z not in (x, y):
+          yield x, y, z
+
+
+def perm3(elems):
+  return list(_perm3(elems))
+
+
+def _perm4(elems):
+  for x in elems:
+    for y in elems:
+      if y == x:
+        continue
+      for z in elems:
+        if z in (x, y):
+          continue
+        for t in elems:
+          if t not in (x, y, z):
+            yield x, y, z, t
+
+
+def perm4(elems):
+  return list(_perm4(elems))

ag4masses/alphageometry/inspect_defs.py

@@ -0,0 +1,27 @@
+import problem as pr
+
+def inspect_definition():
+    # Load definitions from the file
+    defs = pr.Definition.from_txt_file('defs.txt', to_dict=True)
+    
+    # Access the 'semicircle' definition
+    semicircle_def = defs.get('semicircle')
+    
+    if semicircle_def:
+        # Print out the details of the 'semicircle' definition
+        print("Semicircle Definition:")
+        print(semicircle_def)
+        
+        # Print specific attributes of the 'semicircle' definition
+        # Replace 'attribute_name' with the actual attribute names you want to print
+        if hasattr(semicircle_def, 'name'):
+            print(f"Name: {semicircle_def.name}")
+        if hasattr(semicircle_def, 'description'):
+            print(f"Description: {semicircle_def.description}")
+        if hasattr(semicircle_def, 'some_other_attribute'):
+            print(f"Some Other Attribute: {semicircle_def.some_other_attribute}")
+    else:
+        print("No definition found for 'semicircle'")
+
+if __name__ == "__main__":
+    inspect_definition()

ag4masses/alphageometry/lm_inference.py

@@ -1,189 +1,189 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Wrapper for language modeling inference implemented in Meliad."""
-from typing import Any, Dict
-
-import jax
-import models  # pylint: disable=unused-import
-import t5.data
-from transformer import inference_utils
-
-
-np = jax.numpy
-
-
-Trainer = inference_utils.Trainer
-
-MetricsOutput = Dict[str, Any]  # Metrics output by model.
-
-
-parse_gin_configuration = inference_utils.parse_gin_configuration
-
-
-class LanguageModelInference:
-  """Meliad wrapper for LM inference."""
-
-  def __init__(self, vocab_path: str, load_dir: str, mode='beam_search'):
-    self.vocab = t5.data.SentencePieceVocabulary(vocab_path)
-
-    # This task won't be pulling from a dataset.
-    def null_iter_fn() -> None:
-      return None
-
-    process_summaries_f = inference_utils.models.process_summaries_function(
-        self.vocab
-    )
-
-    trainer = inference_utils.training_loop.Trainer(
-        get_training_dataset_iterator=null_iter_fn,
-        get_test_dataset_iterator=None,
-        pretty_print_input_function=None,
-        process_summaries_function=process_summaries_f,
-        load_dir=load_dir,
-        workdir='',  # Don't log or save checkpoints.
-        replicate_mode=False,
-    )  # Run on a single device at batch size 1.
-    self.trainer = trainer
-
-    # Create and initialize the model.
-    (tstate, _, imodel, prngs) = trainer.initialize_model()
-    self.imodel = imodel
-    self.batch_size = imodel.task_config.batch_size
-
-    self.n = imodel.num_heads
-    self.h = imodel.head_size
-
-    # Create an inference task.
-    writers = {}
-    self.task = trainer.create_training_task(mode, imodel, prngs, writers)  # pylint: disable=too-many-function-args
-
-    # Register any additional actions.
-    # Actions are cleared first for use with colab.
-    inference_utils.training_loop.clear_interstep_callbacks()
-    inference_utils.training_loop.register_interstep_callbacks()
-    self.tstate = tstate
-
-    # some default parameters.
-    eos = [0] * 1024
-    for idx in self.encode_list(['.', ';']):
-      eos[idx] = 1
-
-    self.eos = np.array(eos, dtype=np.bfloat16)
-    self.mask = jax.numpy.ones([1024], dtype=np.bfloat16)
-
-  def decode(self, ids: list[int]) -> str:
-    return self.vocab.decode(ids)
-
-  def decode_list(self, tokens: list[int]) -> list[str]:
-    return [self.decode([tok]) for tok in tokens]
-
-  def encode(self, inputs_str: str) -> list[int]:
-    return self.vocab.encode(inputs_str)
-
-  def encode_list(self, inputs_strs: list[str]) -> list[int]:
-    result = [self.vocab.encode(x) for x in inputs_strs]
-    assert all([len(x) == 1 for x in result]), [
-        self.decode(x) for x in result if len(x) != 1
-    ]
-    return [x[0] for x in result]
-
-  def call(
-      self,
-      inputs: np.ndarray,
-      dstate: tuple[dict[str, np.ndarray], ...] = None,
-      eos: np.ndarray = None,
-      mask: np.ndarray = None,
-  ) -> MetricsOutput:
-    """Call the meliad model."""
-    batch_size, length = inputs.shape
-    inputs = jax.numpy.pad(inputs, [(0, 0), (0, 1024 - length)])
-
-    if eos is None:
-      eos = self.eos
-    if mask is None:
-      mask = self.mask
-
-    x = {'targets': inputs, 'length': length, 'eos': eos, 'mask': mask}
-
-    if dstate is not None:
-      x['start_of_sequence'] = jax.numpy.array([False] * batch_size)
-    else:
-      dstate = tuple(
-          [{  # this dummy value will never be used.
-              'current_index': np.array([0] * batch_size, dtype=np.int32),
-              'keys': np.zeros(
-                  (batch_size, 2048, self.n, self.h), dtype=np.bfloat16
-              ),
-              'values': np.zeros(
-                  (batch_size, 2048, self.n, self.h), dtype=np.bfloat16
-              ),
-              'recurrent_kvq': None,
-              'relative_position_bias': np.zeros(
-                  (batch_size, self.n, 1, 1024), dtype=np.bfloat16
-              ),
-          }]
-          * 12
-      )
-      x['start_of_sequence'] = jax.numpy.array([True] * batch_size)
-
-    x['dstate'] = dstate
-    _, metrics_np = self.task.run_step(self.tstate, x, 0)
-    return metrics_np
-
-  def beam_decode(
-      self,
-      inputs: str,
-      eos_tokens: np.ndarray = None,
-      mask_tokens: np.ndarray = None,
-      dstate: dict[str, np.ndarray] = None,
-  ) -> MetricsOutput:
-    """Beam search."""
-    inputs = jax.numpy.array([self.vocab.encode(inputs)] * self.batch_size)
-
-    eos = self.eos
-    if eos_tokens is not None:
-      eos_ids = self.encode_list(eos_tokens)
-      eos = np.array(
-          [1 if idx in eos_ids else 0 for idx in range(1024)], dtype=np.bfloat16
-      ).reshape((1, 1, 1024))
-
-    mask = self.mask
-    if mask_tokens is not None:
-      mask_ids = self.encode_list(mask_tokens)
-      mask = np.array(
-          [0 if idx in mask_ids else 1 for idx in range(1024)],
-          dtype=np.bfloat16,
-      ).reshape((1, 1, 1024))
-
-    metrics_np = self.call(inputs, dstate=dstate, eos=eos, mask=mask)
-
-    finished_seqs = metrics_np['finished_seqs']
-    finished_scores = metrics_np['finished_scores']
-
-    seqs = []
-    scores = []
-    for seq, score in zip(finished_seqs, finished_scores):
-      seq = self.decode(seq[1:])
-      seqs.append(seq)
-      scores.append(score)
-
-    return {
-        'finished_seqs': finished_seqs,
-        'finished_scores': finished_scores,
-        'seqs_str': seqs,
-        'scores': scores,
-        'dstate': metrics_np['dstate'],
-    }
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Wrapper for language modeling inference implemented in Meliad."""
+from typing import Any, Dict
+
+import jax
+import models  # pylint: disable=unused-import
+import t5.data
+from meliad_lib.meliad.transformer import inference_utils
+
+
+np = jax.numpy
+
+
+Trainer = inference_utils.Trainer
+
+MetricsOutput = Dict[str, Any]  # Metrics output by model.
+
+
+parse_gin_configuration = inference_utils.parse_gin_configuration
+
+
+class LanguageModelInference:
+  """Meliad wrapper for LM inference."""
+
+  def __init__(self, vocab_path: str, load_dir: str, mode='beam_search'):
+    self.vocab = t5.data.SentencePieceVocabulary(vocab_path)
+
+    # This task won't be pulling from a dataset.
+    def null_iter_fn() -> None:
+      return None
+
+    process_summaries_f = inference_utils.models.process_summaries_function(
+        self.vocab
+    )
+
+    trainer = inference_utils.training_loop.Trainer(
+        get_training_dataset_iterator=null_iter_fn,
+        get_test_dataset_iterator=None,
+        pretty_print_input_function=None,
+        process_summaries_function=process_summaries_f,
+        load_dir=load_dir,
+        workdir='',  # Don't log or save checkpoints.
+        replicate_mode=False,
+    )  # Run on a single device at batch size 1.
+    self.trainer = trainer
+
+    # Create and initialize the model.
+    (tstate, _, imodel, prngs) = trainer.initialize_model()
+    self.imodel = imodel
+    self.batch_size = imodel.task_config.batch_size
+
+    self.n = imodel.num_heads
+    self.h = imodel.head_size
+
+    # Create an inference task.
+    writers = {}
+    self.task = trainer.create_training_task(mode, imodel, prngs, writers)  # pylint: disable=too-many-function-args
+
+    # Register any additional actions.
+    # Actions are cleared first for use with colab.
+    inference_utils.training_loop.clear_interstep_callbacks()
+    inference_utils.training_loop.register_interstep_callbacks()
+    self.tstate = tstate
+
+    # some default parameters.
+    eos = [0] * 1024
+    for idx in self.encode_list(['.', ';']):
+      eos[idx] = 1
+
+    self.eos = np.array(eos, dtype=np.bfloat16)
+    self.mask = jax.numpy.ones([1024], dtype=np.bfloat16)
+
+  def decode(self, ids: list[int]) -> str:
+    return self.vocab.decode(ids)
+
+  def decode_list(self, tokens: list[int]) -> list[str]:
+    return [self.decode([tok]) for tok in tokens]
+
+  def encode(self, inputs_str: str) -> list[int]:
+    return self.vocab.encode(inputs_str)
+
+  def encode_list(self, inputs_strs: list[str]) -> list[int]:
+    result = [self.vocab.encode(x) for x in inputs_strs]
+    assert all([len(x) == 1 for x in result]), [
+        self.decode(x) for x in result if len(x) != 1
+    ]
+    return [x[0] for x in result]
+
+  def call(
+      self,
+      inputs: np.ndarray,
+      dstate: tuple[dict[str, np.ndarray], ...] = None,
+      eos: np.ndarray = None,
+      mask: np.ndarray = None,
+  ) -> MetricsOutput:
+    """Call the meliad model."""
+    batch_size, length = inputs.shape
+    inputs = jax.numpy.pad(inputs, [(0, 0), (0, 1024 - length)])
+
+    if eos is None:
+      eos = self.eos
+    if mask is None:
+      mask = self.mask
+
+    x = {'targets': inputs, 'length': length, 'eos': eos, 'mask': mask}
+
+    if dstate is not None:
+      x['start_of_sequence'] = jax.numpy.array([False] * batch_size)
+    else:
+      dstate = tuple(
+          [{  # this dummy value will never be used.
+              'current_index': np.array([0] * batch_size, dtype=np.int32),
+              'keys': np.zeros(
+                  (batch_size, 2048, self.n, self.h), dtype=np.bfloat16
+              ),
+              'values': np.zeros(
+                  (batch_size, 2048, self.n, self.h), dtype=np.bfloat16
+              ),
+              'recurrent_kvq': None,
+              'relative_position_bias': np.zeros(
+                  (batch_size, self.n, 1, 1024), dtype=np.bfloat16
+              ),
+          }]
+          * 12
+      )
+      x['start_of_sequence'] = jax.numpy.array([True] * batch_size)
+
+    x['dstate'] = dstate
+    _, metrics_np = self.task.run_step(self.tstate, x, 0)
+    return metrics_np
+
+  def beam_decode(
+      self,
+      inputs: str,
+      eos_tokens: np.ndarray = None,
+      mask_tokens: np.ndarray = None,
+      dstate: dict[str, np.ndarray] = None,
+  ) -> MetricsOutput:
+    """Beam search."""
+    inputs = jax.numpy.array([self.vocab.encode(inputs)] * self.batch_size)
+
+    eos = self.eos
+    if eos_tokens is not None:
+      eos_ids = self.encode_list(eos_tokens)
+      eos = np.array(
+          [1 if idx in eos_ids else 0 for idx in range(1024)], dtype=np.bfloat16
+      ).reshape((1, 1, 1024))
+
+    mask = self.mask
+    if mask_tokens is not None:
+      mask_ids = self.encode_list(mask_tokens)
+      mask = np.array(
+          [0 if idx in mask_ids else 1 for idx in range(1024)],
+          dtype=np.bfloat16,
+      ).reshape((1, 1, 1024))
+
+    metrics_np = self.call(inputs, dstate=dstate, eos=eos, mask=mask)
+
+    finished_seqs = metrics_np['finished_seqs']
+    finished_scores = metrics_np['finished_scores']
+
+    seqs = []
+    scores = []
+    for seq, score in zip(finished_seqs, finished_scores):
+      seq = self.decode(seq[1:])
+      seqs.append(seq)
+      scores.append(score)
+
+    return {
+        'finished_seqs': finished_seqs,
+        'finished_scores': finished_scores,
+        'seqs_str': seqs,
+        'scores': scores,
+        'dstate': metrics_np['dstate'],
+    }

ag4masses/alphageometry/models.py

@@ -1,178 +1,178 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Transformer language model generate mode."""
-
-from typing import Any, Tuple
-import beam_search
-import decoder_stack
-import gin
-import jax
-import jax.numpy as jnp
-from transformer import models
-
-
-@gin.configurable
-class DecoderOnlyLanguageModelGenerate(models.DecoderOnlyLanguageModel):
-  """Decoder only language modeling in inference mode."""
-
-  decoder_factory = decoder_stack.DecoderStackGenerate
-
-  num_heads: int = gin.REQUIRED
-  head_size: int = gin.REQUIRED
-
-  def get_fake_input(self) -> dict[str, Any]:
-    fake_input_dict = super().get_fake_input()
-    b = self.task_config.batch_size
-    n = self.num_heads
-    h = self.head_size
-    fake_input_dict.update({
-        'dstate': tuple(
-            [{
-                'current_index': jnp.array([0] * b, dtype=jnp.int32),
-                'keys': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
-                'values': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
-                'recurrent_kvq': None,
-                'relative_position_bias': jnp.zeros(
-                    (b, n, 1, 1024), dtype=jnp.bfloat16
-                ),
-            }]
-            * 12
-        ),
-        'eos': jnp.zeros([1024], dtype=jnp.bfloat16),
-        'mask': jnp.ones([1024], dtype=jnp.bfloat16),
-        'length': 1,
-        'temperature': 1.0,
-    })
-    return fake_input_dict
-
-  def __call__(self, inputs: ...) -> tuple[Any, dict[str, Any]]:
-    # Make sure this code is not used on untested cases.
-    if self.mode not in ['init', 'beam_search']:
-      raise ValueError(f'{type(self)} cannot do mode {self.mode}')
-    if self.decoder.supports_generate():
-      raise ValueError(f'{type(self)}.decoder cannot supports_generate()')
-
-    self.decoder(
-        input_tokens=inputs['targets'][:, 0:1],
-        target_tokens=None,
-        start_of_sequence=inputs['start_of_sequence'],
-    )
-
-    b = inputs['targets'].shape[0]
-    no_start_of_seq = jnp.array([False] * b, dtype=jnp.bool_)
-
-    # This fn is used in both beam_search or topk_sampling.
-    def tokens_to_logits_fn(
-        input_token: jnp.ndarray, dstate: tuple[dict[str, jnp.ndarray], ...]
-    ) -> tuple[jnp.ndarray, tuple[dict[str, jnp.ndarray], ...]]:
-      (logits, dstate, _) = self.decoder(
-          input_tokens=input_token,
-          target_tokens=None,
-          start_of_sequence=no_start_of_seq,
-          decoder_state=dstate,
-      )
-      return logits[:, -1, :], dstate
-
-    last_token = jax.lax.dynamic_slice_in_dim(
-        inputs['targets'], inputs['length'] - 1, 1, axis=1
-    )
-
-    # last token is used to seed beam_search
-    inputs['targets'] = inputs['targets'][:, 0:-1]
-    dstate = jax.lax.cond(
-        inputs['start_of_sequence'][0],
-        lambda: self.generate(inputs)[0],
-        lambda: inputs['dstate'],
-    )
-
-    # Then we run beam search, init with last_token & dstate.
-    finished_seqs, finished_scores, dstate = beam_search.beam_search_flat(
-        last_token,
-        dstate,
-        tokens_to_logits_fn,
-        max_decode_len=512,
-        eos=inputs['eos'].reshape((1, 1, -1)),
-        mask=inputs['mask'].reshape((1, 1, -1)),
-    )
-
-    return 0.0, {
-        'finished_seqs': finished_seqs,
-        'finished_scores': finished_scores,
-        'dstate': dstate,
-    }
-
-  def generate(
-      self, inputs: ...
-  ) -> tuple[tuple[dict[str, jnp.ndarray, ...], ...], jnp.ndarray]:
-    """Generate an output sequence.
-
-    Args:
-      inputs: the same as argument to _call_.
-
-    Returns:
-      An array of generated tokens of shape (batch_size, sequence_length).
-    """
-    input_tokens = inputs['targets']  # [b,seq_len]
-    start_of_sequence = inputs['start_of_sequence']  # [b]
-    target_tokens = jnp.pad(input_tokens[:, 1:], [(0, 0), (0, 1)])
-    batch_size = target_tokens.shape[0]
-
-    # Assuming all sequences start at the same time.
-    start0 = inputs['start_of_sequence'][0]
-    dstate = jax.lax.cond(
-        start0,
-        lambda: self.decoder.init_decoder_state_vanilla(  # pylint: disable=g-long-lambda
-            1024, start_of_sequence
-        ),
-        lambda: inputs['dstate'],
-    )
-
-    first_token = input_tokens[:, 0:1]
-    no_start_of_seq = jnp.array([False] * batch_size, dtype=jnp.bool_)
-    temperature = 1
-    if 'temperature' in inputs:
-      temperature = inputs['temperature']
-
-    num_steps = inputs['length']
-    if self.mode == 'beam_search':
-      num_steps -= 1
-
-    def cond_fn(scan_state) -> jnp.bool_:
-      _, _, i, _ = scan_state
-      return i < num_steps
-
-    def loop_fn(scan_state: Any) -> Tuple[Any, Any, Any, Any]:
-      (dstate, input_token, i, _) = scan_state
-
-      (logits, dstate, _) = self.decoder(
-          input_tokens=input_token,
-          target_tokens=None,
-          start_of_sequence=no_start_of_seq,
-          decoder_state=dstate,
-      )
-
-      logits = logits / temperature
-      output_token = jax.lax.dynamic_slice_in_dim(target_tokens, i, 1, axis=1)
-
-      return (dstate, output_token, i + 1, logits)
-
-    # Scan over the sequence length.
-    dummy_logits = jnp.zeros((batch_size, 1, 1024))
-    initial_scan_state = (dstate, first_token, 0, dummy_logits)
-    dstate, _, _, logits = jax.lax.while_loop(
-        cond_fn, loop_fn, initial_scan_state
-    )
-    return dstate, logits
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Transformer language model generate mode."""
+
+from typing import Any, Tuple
+import beam_search
+import decoder_stack
+import gin
+import jax
+import jax.numpy as jnp
+from meliad_lib.meliad.transformer import models
+
+
+@gin.configurable
+class DecoderOnlyLanguageModelGenerate(models.DecoderOnlyLanguageModel):
+  """Decoder only language modeling in inference mode."""
+
+  decoder_factory = decoder_stack.DecoderStackGenerate
+
+  num_heads: int = gin.REQUIRED
+  head_size: int = gin.REQUIRED
+
+  def get_fake_input(self) -> dict[str, Any]:
+    fake_input_dict = super().get_fake_input()
+    b = self.task_config.batch_size
+    n = self.num_heads
+    h = self.head_size
+    fake_input_dict.update({
+        'dstate': tuple(
+            [{
+                'current_index': jnp.array([0] * b, dtype=jnp.int32),
+                'keys': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
+                'values': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
+                'recurrent_kvq': None,
+                'relative_position_bias': jnp.zeros(
+                    (b, n, 1, 1024), dtype=jnp.bfloat16
+                ),
+            }]
+            * 12
+        ),
+        'eos': jnp.zeros([1024], dtype=jnp.bfloat16),
+        'mask': jnp.ones([1024], dtype=jnp.bfloat16),
+        'length': 1,
+        'temperature': 1.0,
+    })
+    return fake_input_dict
+
+  def __call__(self, inputs: ...) -> tuple[Any, dict[str, Any]]:
+    # Make sure this code is not used on untested cases.
+    if self.mode not in ['init', 'beam_search']:
+      raise ValueError(f'{type(self)} cannot do mode {self.mode}')
+    if self.decoder.supports_generate():
+      raise ValueError(f'{type(self)}.decoder cannot supports_generate()')
+
+    self.decoder(
+        input_tokens=inputs['targets'][:, 0:1],
+        target_tokens=None,
+        start_of_sequence=inputs['start_of_sequence'],
+    )
+
+    b = inputs['targets'].shape[0]
+    no_start_of_seq = jnp.array([False] * b, dtype=jnp.bool_)
+
+    # This fn is used in both beam_search or topk_sampling.
+    def tokens_to_logits_fn(
+        input_token: jnp.ndarray, dstate: tuple[dict[str, jnp.ndarray], ...]
+    ) -> tuple[jnp.ndarray, tuple[dict[str, jnp.ndarray], ...]]:
+      (logits, dstate, _) = self.decoder(
+          input_tokens=input_token,
+          target_tokens=None,
+          start_of_sequence=no_start_of_seq,
+          decoder_state=dstate,
+      )
+      return logits[:, -1, :], dstate
+
+    last_token = jax.lax.dynamic_slice_in_dim(
+        inputs['targets'], inputs['length'] - 1, 1, axis=1
+    )
+
+    # last token is used to seed beam_search
+    inputs['targets'] = inputs['targets'][:, 0:-1]
+    dstate = jax.lax.cond(
+        inputs['start_of_sequence'][0],
+        lambda: self.generate(inputs)[0],
+        lambda: inputs['dstate'],
+    )
+
+    # Then we run beam search, init with last_token & dstate.
+    finished_seqs, finished_scores, dstate = beam_search.beam_search_flat(
+        last_token,
+        dstate,
+        tokens_to_logits_fn,
+        max_decode_len=512,
+        eos=inputs['eos'].reshape((1, 1, -1)),
+        mask=inputs['mask'].reshape((1, 1, -1)),
+    )
+
+    return 0.0, {
+        'finished_seqs': finished_seqs,
+        'finished_scores': finished_scores,
+        'dstate': dstate,
+    }
+
+  def generate(
+      self, inputs: ...
+  ) -> tuple[tuple[dict[str, jnp.ndarray, ...], ...], jnp.ndarray]:
+    """Generate an output sequence.
+
+    Args:
+      inputs: the same as argument to _call_.
+
+    Returns:
+      An array of generated tokens of shape (batch_size, sequence_length).
+    """
+    input_tokens = inputs['targets']  # [b,seq_len]
+    start_of_sequence = inputs['start_of_sequence']  # [b]
+    target_tokens = jnp.pad(input_tokens[:, 1:], [(0, 0), (0, 1)])
+    batch_size = target_tokens.shape[0]
+
+    # Assuming all sequences start at the same time.
+    start0 = inputs['start_of_sequence'][0]
+    dstate = jax.lax.cond(
+        start0,
+        lambda: self.decoder.init_decoder_state_vanilla(  # pylint: disable=g-long-lambda
+            1024, start_of_sequence
+        ),
+        lambda: inputs['dstate'],
+    )
+
+    first_token = input_tokens[:, 0:1]
+    no_start_of_seq = jnp.array([False] * batch_size, dtype=jnp.bool_)
+    temperature = 1
+    if 'temperature' in inputs:
+      temperature = inputs['temperature']
+
+    num_steps = inputs['length']
+    if self.mode == 'beam_search':
+      num_steps -= 1
+
+    def cond_fn(scan_state) -> jnp.bool_:
+      _, _, i, _ = scan_state
+      return i < num_steps
+
+    def loop_fn(scan_state: Any) -> Tuple[Any, Any, Any, Any]:
+      (dstate, input_token, i, _) = scan_state
+
+      (logits, dstate, _) = self.decoder(
+          input_tokens=input_token,
+          target_tokens=None,
+          start_of_sequence=no_start_of_seq,
+          decoder_state=dstate,
+      )
+
+      logits = logits / temperature
+      output_token = jax.lax.dynamic_slice_in_dim(target_tokens, i, 1, axis=1)
+
+      return (dstate, output_token, i + 1, logits)
+
+    # Scan over the sequence length.
+    dummy_logits = jnp.zeros((batch_size, 1, 1024))
+    initial_scan_state = (dstate, first_token, 0, dummy_logits)
+    dstate, _, _, logits = jax.lax.while_loop(
+        cond_fn, loop_fn, initial_scan_state
+    )
+    return dstate, logits

ag4masses/alphageometry/numericals.py

@@ -25,14 +25,13 @@ from matplotlib import pyplot as plt
 import matplotlib.colors as mcolors
 import numpy as np
 from numpy.random import uniform as unif  # pylint: disable=g-importing-member
+import graph as gh
 
-
-matplotlib.use('Agg')
+matplotlib.use('TkAgg')
 
 
 ATOM = 1e-12
 
-
 # Some variables are there for better code reading.
 # pylint: disable=unused-assignment
 # pylint: disable=unused-argument
@@ -440,6 +439,75 @@ class Circle:
     return [result]
 
 
+class SemiCircle(Circle):
+  """Numerical semicircle, inherits from Circle."""
+
+  def __init__(
+      self,
+      center: Optional[Point] = None,
+      radius: Optional[float] = None,
+      p1: Optional[Point] = None,
+      p2: Optional[Point] = None,
+      p3: Optional[Point] = None,
+  ):
+      self.p1 = p1
+      self.p2 = p2
+      self.p3 = p3
+      # Initialize as a Circle
+      super().__init__(center, radius, p1, p2, p3)
+      # If p1 and p2 define a diameter, set the center and radius accordingly
+      if p1 and p2 and not center:
+          self.center = Point((p1.x + p2.x) / 2, (p1.y + p2.y) / 2)
+          self.radius = p1.distance(p2) / 2
+          self.r2 = self.radius ** 2
+
+      # Define the direction or plane for the semicircle (important for sampling and boundaries)
+
+  def is_within_boundary(self, point: Point) -> bool:
+      """Check if a point is within the boundary of the semicircle."""
+      vector_to_point = point - self.center
+      angle = math.atan2(vector_to_point.y, vector_to_point.x)
+
+      # Normalize the angle within [0, 2*pi]
+      angle = angle if angle >= 0 else (2 * np.pi + angle)
+      
+      # Check if the point is within the semicircle (half of the circle)
+      return -np.pi / 2 <= angle <= np.pi / 2
+
+  def sample_within(self, points: list[Point], n: int = 5) -> list[Point]:
+      """Sample a point within the semicircle."""
+      result = None
+      best = -1.0
+      for _ in range(n):
+          # Generate a random angle between -π/2 and π/2 for the semicircle
+          ang = unif(-0.5, 0.5) * np.pi
+          x = self.center + Point(np.cos(ang), np.sin(ang)) * self.radius
+
+          # Check if the sampled point is within the active part of the semicircle
+          if not self.is_within_boundary(x):
+              continue
+
+          # Find the minimum distance between the generated point and the provided points
+          mind = min([x.distance(p) for p in points])
+          if mind > best:
+              best = mind
+              result = x
+
+      return [result]
+
+  def intersect(self, obj: Union[Line, Circle]) -> tuple[Point, ...]:
+      """Find intersection points with a Line or another Circle, constrained to the semicircle."""
+      if isinstance(obj, Line):
+          intersections = obj.intersect(self)
+      elif isinstance(obj, Circle):
+          intersections = circle_circle_intersection(self, obj)
+      else:
+          return tuple()
+
+      # Filter intersections to only return points within the semicircle
+      return tuple(p for p in intersections if self.is_within_boundary(p))
+
+
 class HoleCircle(Circle):
   """Numerical circle with a missing point."""
 
@@ -565,6 +633,18 @@ def circle_segment_intersect(
     result.append(py)
   return result
 
+def semicircle_segment_intersect(
+    circle: SemiCircle, p1: Point, p2: Point
+) -> list[Point]:
+  l = Line(p1, p2)
+  px, py = line_circle_intersection(l, circle)
+
+  result = []
+  if _check_between(px, p1, p2):
+    result.append(px)
+  if _check_between(py, p1, p2):
+    result.append(py)
+  return result
 
 def line_segment_intersection(l: Line, A: Point, B: Point) -> Point:  # pylint: disable=invalid-name
   a, b, c = l.coefficients
@@ -656,6 +736,13 @@ def check_circle(points: list[Point]) -> bool:
   oa, ob, oc = o.distance(a), o.distance(b), o.distance(c)
   return close_enough(oa, ob) and close_enough(ob, oc)
 
+def check_semicircle(points: list[Point]) -> bool:
+  if len(points) != 4:
+    return False
+  o, a, b, c = points
+  oa, ob, oc = o.distance(a), o.distance(b), o.distance(c)
+  return close_enough(oa, ob) and close_enough(ob, oc)
+
 
 def check_coll(points: list[Point]) -> bool:
   a, b = points[:2]
@@ -894,10 +981,12 @@ def naming_position(
   _ = ax
   r = 0.08
   c = Circle(center=p, radius=r)
+  sc = SemiCircle(center=p, radius=r)
   avoid = []
   for p1, p2 in lines:
     try:
       avoid.extend(circle_segment_intersect(c, p1, p2))
+      avoid.extend(semicircle_segment_intersect(sc, p1, p2))
     except InvalidQuadSolveError:
       continue
   for x in circles:
@@ -928,6 +1017,7 @@ def draw_point(
     name: str,
     lines: list[Line],
     circles: list[Circle],
+    semicircles: list[SemiCircle],
     color: Any = 'white',
     size: float = 15,
 ) -> None:
@@ -1029,6 +1119,133 @@ def draw_circle(
   _draw_circle(ax, circle, color)
   return circle
 
+def check_points_semicircle(p1, p2, p3):
+    """
+    Check if three points are in a semicircle, and determine the circle center, radius,
+    and points forming the diameter if applicable. If no pair forms a diameter, calculate
+    the circle center passing through all three points.
+
+    Parameters:
+        p1, p2, p3 (tuple): Three points as (x, y) coordinates.
+    
+    Returns:
+        dict: A dictionary containing:
+              - 'center': (cx, cy), the circle center.
+              - 'radius': The radius of the circle.
+              - 'diameter_points': A tuple of two points that form the diameter (or None).
+              - 'is_valid': True if a circle can be formed; False otherwise.
+    """
+    # Unpack points
+    x1, y1 = p1
+    x2, y2 = p2
+    x3, y3 = p3
+    
+    # Calculate circumcenter
+    A = np.array([[x1 - x2, y1 - y2], [x1 - x3, y1 - y3]])
+    B = np.array([((x1**2 - x2**2) + (y1**2 - y2**2)) / 2, ((x1**2 - x3**2) + (y1**2 - y3**2)) / 2])
+    
+    try:
+        center = np.linalg.solve(A, B)  # Solving linear system to get circle center
+    except np.linalg.LinAlgError:
+        return {'is_valid': False}  # Points are collinear, no unique circle
+
+    cx, cy = center
+    radius = np.sqrt((x1 - cx)**2 + (y1 - cy)**2)
+
+    # Function to check if two points form a diameter
+    def is_diameter(px, py, qx, qy):
+        midpoint_x, midpoint_y = (px + qx) / 2, (py + qy) / 2
+        return np.isclose(midpoint_x, cx) and np.isclose(midpoint_y, cy)
+    
+    # Check for diameter
+    if is_diameter(x1, y1, x2, y2):
+        diameter_points = (p1, p2)
+    elif is_diameter(x1, y1, x3, y3):
+        diameter_points = (p1, p3)
+    elif is_diameter(x2, y2, x3, y3):
+        diameter_points = (p2, p3)
+    else:
+        diameter_points = None  # No pair forms a diameter; use circumcenter
+    
+    return {
+        'center': center,
+        'radius': radius,
+        'diameter_points': diameter_points,
+        'is_valid': True
+    }
+
+def _draw_semicircle(
+    ax: matplotlib.axes.Axes, P1: Point, P2: Point, P3: Point, color: Any = 'cyan', lw: float = 1.2
+) -> None:
+    """
+    Draws a semicircle passing through three points or with one or two points on the diameter.
+
+    Parameters:
+        ax (matplotlib.axes.Axes): The Matplotlib Axes on which the semicircle will be drawn.
+        P1, P2, P3 (Point): The three points through which the semicircle will pass.
+        color (Any): Color of the semicircle.
+        lw (float): Line width of the semicircle.
+    """
+    result = check_points_semicircle((P1.x, P1.y), (P2.x, P2.y), (P3.x, P3.y))
+    if not result['is_valid']:
+        print("Points are collinear; cannot form a semicircle.")
+        return
+    
+    cx, cy = result['center']
+    radius = result['radius']
+    diameter_points = result['diameter_points']
+
+    # If no pair forms a diameter, determine angles for all three points
+    if diameter_points is None:
+        # Calculate angles of all three points relative to the circle's center
+        angles = np.arctan2(
+            [P1.y - cy, P2.y - cy, P3.y - cy], 
+            [P1.x - cx, P2.x - cx, P3.x - cx]
+        )
+        angles = (angles + 2 * np.pi) % (2 * np.pi)  # Normalize to [0, 2π]
+        
+        # Determine the start and end angle for the semicircle
+        start_angle = np.min(angles)
+        end_angle = np.max(angles)
+        if end_angle - start_angle > np.pi:
+            start_angle, end_angle = end_angle, start_angle + 2 * np.pi
+    else:
+        # Use diameter points to define the semicircle angles
+        px, py = diameter_points[0]
+        qx, qy = diameter_points[1]
+        start_angle = np.arctan2(py - cy, px - cx)
+        end_angle = np.arctan2(qy - cy, qx - cx)
+        if end_angle - start_angle > np.pi:
+            start_angle, end_angle = end_angle, start_angle + 2 * np.pi
+
+    # Generate points for the semicircle
+    t = np.linspace(start_angle, end_angle, 100)
+    x = cx + radius * np.cos(t)
+    y = cy + radius * np.sin(t)
+    
+    # Plot the semicircle
+    ax.plot(x, y, color=color, lw=lw)
+    
+def draw_semicircle(
+    ax: matplotlib.axes.Axes, semicircle: SemiCircle, color: Any = 'cyan'
+) -> SemiCircle:
+    """Draw a semicircle."""
+    if semicircle.num is not None:
+        semicircle = semicircle.num
+    else:
+        points = semicircle.neighbors(gm.Point)
+        if len(points) <= 2:
+            return
+        points = [p.num for p in points]
+        p1, p2, p3 = points[:3]
+        semicircle = SemiCircle(p1=p1, p2=p2, p3=p3)
+    print(semicircle.p1, semicircle.p2, semicircle.p3)
+    _draw_semicircle(ax, semicircle.p1, semicircle.p2, semicircle.p3, color=color)
+    _draw_line(ax, semicircle.p1, semicircle.p2)
+    _draw_line(ax, semicircle.p2, semicircle.p3)
+    _draw_line(ax, semicircle.p1, semicircle.p3)  
+    return semicircle
+
 
 def mark_segment(
     ax: matplotlib.axes.Axes, p1: Point, p2: Point, color: Any, alpha: float
@@ -1126,7 +1343,9 @@ def highlight(
     _draw_line(ax, c, d, color=color2, lw=2.0, alpha=0.5)
     _draw_line(ax, m, n, color=color1, lw=2.0, alpha=0.5)
     _draw_line(ax, p, q, color=color2, lw=2.0, alpha=0.5)
-
+  elif name == 'semicircle':
+    o, a, b, c = args
+    _draw_semicircle(ax, SemiCircle(center=o, p1=a, p2=b, p3=c), color=color1, lw=2.0)
 
 HCOLORS = None
 
@@ -1136,6 +1355,7 @@ def _draw(
     points: list[gm.Point],
     lines: list[gm.Line],
     circles: list[gm.Circle],
+    semicircles: list[gm.SemiCircle],
     goal: Any,
     equals: list[tuple[Any, Any]],
     highlights: list[tuple[str, list[gm.Point]]],
@@ -1158,9 +1378,10 @@ def _draw(
     p1, p2 = draw_line(ax, l, color=lcolor)
     line_boundaries.append((p1, p2))
   circles = [draw_circle(ax, c, color=ccolor) for c in circles]
+  semicircles = [draw_semicircle(ax, c, color=ccolor) for c in semicircles]
 
   for p in points:
-    draw_point(ax, p.num, p.name, line_boundaries, circles, color=pcolor)
+    draw_point(ax, p.num, p.name, line_boundaries, circles, semicircles, color=pcolor)
 
   if equals:
     for i, segs in enumerate(equals['segments']):
@@ -1204,6 +1425,7 @@ def draw(
     points: list[gm.Point],
     lines: list[gm.Line],
     circles: list[gm.Circle],
+    semicircles: list[gm.SemiCircle],
     segments: list[gm.Segment],
     goal: Any = None,
     highlights: list[tuple[str, list[gm.Point]]] = None,
@@ -1214,8 +1436,8 @@ def draw(
 ) -> None:
   """Draw everything on the same canvas."""
   plt.close()
-  imsize = 512 / 100
-  fig, ax = plt.subplots(figsize=(imsize, imsize), dpi=100)
+  imsize = 1280 / 200
+  fig, ax = plt.subplots(figsize=(imsize, imsize), dpi=200)
 
   set_theme(theme)
 
@@ -1224,7 +1446,7 @@ def draw(
   else:
     ax.set_facecolor((1.0, 1.0, 1.0))
 
-  _draw(ax, points, lines, circles, goal, equals, highlights)
+  _draw(ax, points, lines, circles, semicircles, goal, equals, highlights)
 
   plt.axis('equal')
   fig.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=0, hspace=0)
@@ -1238,8 +1460,6 @@ def draw(
     plt.savefig(save_to)
   # plt.show(block=block)
   
-
-
 def close_enough(a: float, b: float, tol: float = 1e-12) -> bool:
   return abs(a - b) < tol
 
@@ -1560,6 +1780,9 @@ def sketch_circle(args: tuple[gm.Point, ...]) -> Circle:
   a, b, c = args
   return Circle(center=a, radius=b.distance(c))
 
+def sketch_semicircle(args: tuple[gm.Point, ...]) -> SemiCircle:
+  a, b, c = args
+  return SemiCircle(center=a, radius=b.distance(c), p1=b, p2=c)
 
 def sketch_cc_tangent(args: tuple[gm.Point, ...]) -> tuple[Point, ...]:
   """Sketch tangents to two circles."""

ag4masses/alphageometry/pretty.py

@@ -1,216 +1,216 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Utilities for string manipulation in the DSL."""
-
-MAP_SYMBOL = {
-    'T': 'perp',
-    'P': 'para',
-    'D': 'cong',
-    'S': 'simtri',
-    'I': 'circle',
-    'M': 'midp',
-    'O': 'cyclic',
-    'C': 'coll',
-    '^': 'eqangle',
-    '/': 'eqratio',
-    '%': 'eqratio',
-    '=': 'contri',
-    'X': 'collx',
-    'A': 'acompute',
-    'R': 'rcompute',
-    'Q': 'fixc',
-    'E': 'fixl',
-    'V': 'fixb',
-    'H': 'fixt',
-    'Z': 'fixp',
-    'Y': 'ind',
-}
-
-
-def map_symbol(c: str) -> str:
-  return MAP_SYMBOL[c]
-
-
-def map_symbol_inv(c: str) -> str:
-  return {v: k for k, v in MAP_SYMBOL.items()}[c]
-
-
-def _gcd(x: int, y: int) -> int:
-  while y:
-    x, y = y, x % y
-  return x
-
-
-def simplify(n: int, d: int) -> tuple[int, int]:
-  g = _gcd(n, d)
-  return (n // g, d // g)
-
-
-def pretty2r(a: str, b: str, c: str, d: str) -> str:
-  if b in (c, d):
-    a, b = b, a
-
-  if a == d:
-    c, d = d, c
-
-  return f'{a} {b} {c} {d}'
-
-
-def pretty2a(a: str, b: str, c: str, d: str) -> str:
-  if b in (c, d):
-    a, b = b, a
-
-  if a == d:
-    c, d = d, c
-
-  return f'{a} {b} {c} {d}'
-
-
-def pretty_angle(a: str, b: str, c: str, d: str) -> str:
-  if b in (c, d):
-    a, b = b, a
-  if a == d:
-    c, d = d, c
-
-  if a == c:
-    return f'\u2220{b}{a}{d}'
-  return f'\u2220({a}{b}-{c}{d})'
-
-
-def pretty_nl(name: str, args: list[str]) -> str:
-  """Natural lang formatting a predicate."""
-  if name == 'aconst':
-    a, b, c, d, y = args
-    return f'{pretty_angle(a, b, c, d)} = {y}'
-  if name == 'rconst':
-    a, b, c, d, y = args
-    return f'{a}{b}:{c}{d} = {y}'
-  if name == 'acompute':
-    a, b, c, d = args
-    return f'{pretty_angle(a, b, c, d)}'
-  if name in ['coll', 'C']:
-    return '' + ','.join(args) + ' are collinear'
-  if name == 'collx':
-    return '' + ','.join(list(set(args))) + ' are collinear'
-  if name in ['cyclic', 'O']:
-    return '' + ','.join(args) + ' are concyclic'
-  if name in ['midp', 'midpoint', 'M']:
-    x, a, b = args
-    return f'{x} is midpoint of {a}{b}'
-  if name in ['eqangle', 'eqangle6', '^']:
-    a, b, c, d, e, f, g, h = args
-    return f'{pretty_angle(a, b, c, d)} = {pretty_angle(e, f, g, h)}'
-  if name in ['eqratio', 'eqratio6', '/']:
-    return '{}{}:{}{} = {}{}:{}{}'.format(*args)
-  if name == 'eqratio3':
-    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
-    return f'S {o} {a} {b} {o} {c} {d}'
-  if name in ['cong', 'D']:
-    a, b, c, d = args
-    return f'{a}{b} = {c}{d}'
-  if name in ['perp', 'T']:
-    if len(args) == 2:  # this is algebraic derivation.
-      ab, cd = args  # ab = 'd( ... )'
-      return f'{ab} \u27c2 {cd}'
-    a, b, c, d = args
-    return f'{a}{b} \u27c2 {c}{d}'
-  if name in ['para', 'P']:
-    if len(args) == 2:  # this is algebraic derivation.
-      ab, cd = args  # ab = 'd( ... )'
-      return f'{ab} \u2225 {cd}'
-    a, b, c, d = args
-    return f'{a}{b} \u2225 {c}{d}'
-  if name in ['simtri2', 'simtri', 'simtri*']:
-    a, b, c, x, y, z = args
-    return f'\u0394{a}{b}{c} is similar to \u0394{x}{y}{z}'
-  if name in ['contri2', 'contri', 'contri*']:
-    a, b, c, x, y, z = args
-    return f'\u0394{a}{b}{c} is congruent to \u0394{x}{y}{z}'
-  if name in ['circle', 'I']:
-    o, a, b, c = args
-    return f'{o} is the circumcenter of \\Delta {a}{b}{c}'
-  if name == 'foot':
-    a, b, c, d = args
-    return f'{a} is the foot of {b} on {c}{d}'
-
-
-def pretty(txt: str) -> str:
-  """Pretty formating a predicate string."""
-  if isinstance(txt, str):
-    txt = txt.split(' ')
-  name, *args = txt
-  if name == 'ind':
-    return 'Y ' + ' '.join(args)
-  if name in ['fixc', 'fixl', 'fixb', 'fixt', 'fixp']:
-    return map_symbol_inv(name) + ' ' + ' '.join(args)
-  if name == 'acompute':
-    a, b, c, d = args
-    return 'A ' + ' '.join(args)
-  if name == 'rcompute':
-    a, b, c, d = args
-    return 'R ' + ' '.join(args)
-  if name == 'aconst':
-    a, b, c, d, y = args
-    return f'^ {pretty2a(a, b, c, d)} {y}'
-  if name == 'rconst':
-    a, b, c, d, y = args
-    return f'/ {pretty2r(a, b, c, d)} {y}'
-  if name == 'coll':
-    return 'C ' + ' '.join(args)
-  if name == 'collx':
-    return 'X ' + ' '.join(args)
-  if name == 'cyclic':
-    return 'O ' + ' '.join(args)
-  if name in ['midp', 'midpoint']:
-    x, a, b = args
-    return f'M {x} {a} {b}'
-  if name == 'eqangle':
-    a, b, c, d, e, f, g, h = args
-    return f'^ {pretty2a(a, b, c, d)} {pretty2a(e, f, g, h)}'
-  if name == 'eqratio':
-    a, b, c, d, e, f, g, h = args
-    return f'/ {pretty2r(a, b, c, d)} {pretty2r(e, f, g, h)}'
-  if name == 'eqratio3':
-    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
-    return f'S {o} {a} {b} {o} {c} {d}'
-  if name == 'cong':
-    a, b, c, d = args
-    return f'D {a} {b} {c} {d}'
-  if name == 'perp':
-    if len(args) == 2:  # this is algebraic derivation.
-      ab, cd = args  # ab = 'd( ... )'
-      return f'T {ab} {cd}'
-    a, b, c, d = args
-    return f'T {a} {b} {c} {d}'
-  if name == 'para':
-    if len(args) == 2:  # this is algebraic derivation.
-      ab, cd = args  # ab = 'd( ... )'
-      return f'P {ab} {cd}'
-    a, b, c, d = args
-    return f'P {a} {b} {c} {d}'
-  if name in ['simtri2', 'simtri', 'simtri*']:
-    a, b, c, x, y, z = args
-    return f'S {a} {b} {c} {x} {y} {z}'
-  if name in ['contri2', 'contri', 'contri*']:
-    a, b, c, x, y, z = args
-    return f'= {a} {b} {c} {x} {y} {z}'
-  if name == 'circle':
-    o, a, b, c = args
-    return f'I {o} {a} {b} {c}'
-  if name == 'foot':
-    a, b, c, d = args
-    return f'F {a} {b} {c} {d}'
-  return ' '.join(txt)
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utilities for string manipulation in the DSL."""
+
+MAP_SYMBOL = {
+    'T': 'perp',
+    'P': 'para',
+    'D': 'cong',
+    'S': 'simtri',
+    'I': 'circle',
+    'M': 'midp',
+    'O': 'cyclic',
+    'C': 'coll',
+    '^': 'eqangle',
+    '/': 'eqratio',
+    '%': 'eqratio',
+    '=': 'contri',
+    'X': 'collx',
+    'A': 'acompute',
+    'R': 'rcompute',
+    'Q': 'fixc',
+    'E': 'fixl',
+    'V': 'fixb',
+    'H': 'fixt',
+    'Z': 'fixp',
+    'Y': 'ind',
+}
+
+
+def map_symbol(c: str) -> str:
+  return MAP_SYMBOL[c]
+
+
+def map_symbol_inv(c: str) -> str:
+  return {v: k for k, v in MAP_SYMBOL.items()}[c]
+
+
+def _gcd(x: int, y: int) -> int:
+  while y:
+    x, y = y, x % y
+  return x
+
+
+def simplify(n: int, d: int) -> tuple[int, int]:
+  g = _gcd(n, d)
+  return (n // g, d // g)
+
+
+def pretty2r(a: str, b: str, c: str, d: str) -> str:
+  if b in (c, d):
+    a, b = b, a
+
+  if a == d:
+    c, d = d, c
+
+  return f'{a} {b} {c} {d}'
+
+
+def pretty2a(a: str, b: str, c: str, d: str) -> str:
+  if b in (c, d):
+    a, b = b, a
+
+  if a == d:
+    c, d = d, c
+
+  return f'{a} {b} {c} {d}'
+
+
+def pretty_angle(a: str, b: str, c: str, d: str) -> str:
+  if b in (c, d):
+    a, b = b, a
+  if a == d:
+    c, d = d, c
+
+  if a == c:
+    return f'\u2220{b}{a}{d}'
+  return f'\u2220({a}{b}-{c}{d})'
+
+
+def pretty_nl(name: str, args: list[str]) -> str:
+  """Natural lang formatting a predicate."""
+  if name == 'aconst':
+    a, b, c, d, y = args
+    return f'{pretty_angle(a, b, c, d)} = {y}'
+  if name == 'rconst':
+    a, b, c, d, y = args
+    return f'{a}{b}:{c}{d} = {y}'
+  if name == 'acompute':
+    a, b, c, d = args
+    return f'{pretty_angle(a, b, c, d)}'
+  if name in ['coll', 'C']:
+    return '' + ','.join(args) + ' are collinear'
+  if name == 'collx':
+    return '' + ','.join(list(set(args))) + ' are collinear'
+  if name in ['cyclic', 'O']:
+    return '' + ','.join(args) + ' are concyclic'
+  if name in ['midp', 'midpoint', 'M']:
+    x, a, b = args
+    return f'{x} is midpoint of {a}{b}'
+  if name in ['eqangle', 'eqangle6', '^']:
+    a, b, c, d, e, f, g, h = args
+    return f'{pretty_angle(a, b, c, d)} = {pretty_angle(e, f, g, h)}'
+  if name in ['eqratio', 'eqratio6', '/']:
+    return '{}{}:{}{} = {}{}:{}{}'.format(*args)
+  if name == 'eqratio3':
+    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
+    return f'S {o} {a} {b} {o} {c} {d}'
+  if name in ['cong', 'D']:
+    a, b, c, d = args
+    return f'{a}{b} = {c}{d}'
+  if name in ['perp', 'T']:
+    if len(args) == 2:  # this is algebraic derivation.
+      ab, cd = args  # ab = 'd( ... )'
+      return f'{ab} \u27c2 {cd}'
+    a, b, c, d = args
+    return f'{a}{b} \u27c2 {c}{d}'
+  if name in ['para', 'P']:
+    if len(args) == 2:  # this is algebraic derivation.
+      ab, cd = args  # ab = 'd( ... )'
+      return f'{ab} \u2225 {cd}'
+    a, b, c, d = args
+    return f'{a}{b} \u2225 {c}{d}'
+  if name in ['simtri2', 'simtri', 'simtri*']:
+    a, b, c, x, y, z = args
+    return f'\u0394{a}{b}{c} is similar to \u0394{x}{y}{z}'
+  if name in ['contri2', 'contri', 'contri*']:
+    a, b, c, x, y, z = args
+    return f'\u0394{a}{b}{c} is congruent to \u0394{x}{y}{z}'
+  if name in ['circle', 'I']:
+    o, a, b, c = args
+    return f'{o} is the circumcenter of \\Delta {a}{b}{c}'
+  if name == 'foot':
+    a, b, c, d = args
+    return f'{a} is the foot of {b} on {c}{d}'
+
+
+def pretty(txt: str) -> str:
+  """Pretty formating a predicate string."""
+  if isinstance(txt, str):
+    txt = txt.split(' ')
+  name, *args = txt
+  if name == 'ind':
+    return 'Y ' + ' '.join(args)
+  if name in ['fixc', 'fixl', 'fixb', 'fixt', 'fixp']:
+    return map_symbol_inv(name) + ' ' + ' '.join(args)
+  if name == 'acompute':
+    a, b, c, d = args
+    return 'A ' + ' '.join(args)
+  if name == 'rcompute':
+    a, b, c, d = args
+    return 'R ' + ' '.join(args)
+  if name == 'aconst':
+    a, b, c, d, y = args
+    return f'^ {pretty2a(a, b, c, d)} {y}'
+  if name == 'rconst':
+    a, b, c, d, y = args
+    return f'/ {pretty2r(a, b, c, d)} {y}'
+  if name == 'coll':
+    return 'C ' + ' '.join(args)
+  if name == 'collx':
+    return 'X ' + ' '.join(args)
+  if name == 'cyclic':
+    return 'O ' + ' '.join(args)
+  if name in ['midp', 'midpoint']:
+    x, a, b = args
+    return f'M {x} {a} {b}'
+  if name == 'eqangle':
+    a, b, c, d, e, f, g, h = args
+    return f'^ {pretty2a(a, b, c, d)} {pretty2a(e, f, g, h)}'
+  if name == 'eqratio':
+    a, b, c, d, e, f, g, h = args
+    return f'/ {pretty2r(a, b, c, d)} {pretty2r(e, f, g, h)}'
+  if name == 'eqratio3':
+    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
+    return f'S {o} {a} {b} {o} {c} {d}'
+  if name == 'cong':
+    a, b, c, d = args
+    return f'D {a} {b} {c} {d}'
+  if name == 'perp':
+    if len(args) == 2:  # this is algebraic derivation.
+      ab, cd = args  # ab = 'd( ... )'
+      return f'T {ab} {cd}'
+    a, b, c, d = args
+    return f'T {a} {b} {c} {d}'
+  if name == 'para':
+    if len(args) == 2:  # this is algebraic derivation.
+      ab, cd = args  # ab = 'd( ... )'
+      return f'P {ab} {cd}'
+    a, b, c, d = args
+    return f'P {a} {b} {c} {d}'
+  if name in ['simtri2', 'simtri', 'simtri*']:
+    a, b, c, x, y, z = args
+    return f'S {a} {b} {c} {x} {y} {z}'
+  if name in ['contri2', 'contri', 'contri*']:
+    a, b, c, x, y, z = args
+    return f'= {a} {b} {c} {x} {y} {z}'
+  if name == 'circle':
+    o, a, b, c = args
+    return f'I {o} {a} {b} {c}'
+  if name == 'foot':
+    a, b, c, d = args
+    return f'F {a} {b} {c} {d}'
+  return ' '.join(txt)

ag4masses/alphageometry/problem.py

@@ -1,1133 +1,1152 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Implements objects to represent problems, theorems, proofs, traceback."""
-
-from __future__ import annotations
-
-from collections import defaultdict  # pylint: disable=g-importing-member
-from typing import Any
-
-import geometry as gm
-import pretty as pt
-
-
-# pylint: disable=protected-access
-# pylint: disable=unused-variable
-# pylint: disable=unused-argument
-# pylint: disable=unused-assignment
-
-
-def reshape(l: list[Any], n: int = 1) -> list[list[Any]]:
-  assert len(l) % n == 0
-  columns = [[] for i in range(n)]
-  for i, x in enumerate(l):
-    columns[i % n].append(x)
-  return zip(*columns)
-
-
-def isint(x: str) -> bool:
-  try:
-    int(x)
-    return True
-  except:  # pylint: disable=bare-except
-    return False
-
-
-class Construction:
-  """One predicate."""
-
-  @classmethod
-  def from_txt(cls, data: str) -> Construction:
-    data = data.split(' ')
-    return Construction(data[0], data[1:])
-
-  def __init__(self, name: str, args: list[str]):
-    self.name = name
-    self.args = args
-
-  def translate(self, mapping: dict[str, str]) -> Construction:
-    args = [a if isint(a) else mapping[a] for a in self.args]
-    return Construction(self.name, args)
-
-  def txt(self) -> str:
-    return ' '.join([self.name] + list(self.args))
-
-
-class Clause:
-  """One construction (>= 1 predicate)."""
-
-  @classmethod
-  def from_txt(cls, data: str) -> Clause:
-    if data == ' =':
-      return Clause([], [])
-    points, constructions = data.split(' = ')
-    return Clause(
-        points.split(' '),
-        [Construction.from_txt(c) for c in constructions.split(', ')],
-    )
-
-  def __init__(self, points: list[str], constructions: list[Construction]):
-    self.points = []
-    self.nums = []
-
-    for p in points:
-      num = None
-      if isinstance(p, str) and '@' in p:
-        p, num = p.split('@')
-        x, y = num.split('_')
-        num = float(x), float(y)
-      self.points.append(p)
-      self.nums.append(num)
-
-    self.constructions = constructions
-
-  def translate(self, mapping: dict[str, str]) -> Clause:
-    points0 = []
-    for p in self.points:
-      pcount = len(mapping) + 1
-      name = chr(96 + pcount)
-      if name > 'z':  # pcount = 26 -> name = 'z'
-        name = chr(97 + (pcount - 1) % 26) + str((pcount - 1) // 26)
-
-      p0 = mapping.get(p, name)
-      mapping[p] = p0
-      points0.append(p0)
-    return Clause(points0, [c.translate(mapping) for c in self.constructions])
-
-  def add(self, name: str, args: list[str]) -> None:
-    self.constructions.append(Construction(name, args))
-
-  def txt(self) -> str:
-    return (
-        ' '.join(self.points)
-        + ' = '
-        + ', '.join(c.txt() for c in self.constructions)
-    )
-
-
-def _gcd(x: int, y: int) -> int:
-  while y:
-    x, y = y, x % y
-  return x
-
-
-def simplify(n: int, d: int) -> tuple[int, int]:
-  g = _gcd(n, d)
-  return (n // g, d // g)
-
-
-def compare_fn(dep: Dependency) -> tuple[Dependency, str]:
-  return (dep, pt.pretty(dep))
-
-
-def sort_deps(deps: list[Dependency]) -> list[Dependency]:
-  return sorted(deps, key=compare_fn)
-
-
-class Problem:
-  """Describe one problem to solve."""
-
-  @classmethod
-  def from_txt_file(
-      cls, fname: str, to_dict: bool = False, translate: bool = True
-  ):
-    """Load a problem from a text file."""
-    with open(fname, 'r') as f:
-      lines = f.read().split('\n')
-
-    lines = [l for l in lines if l]
-    data = [
-        cls.from_txt(url + '\n' + problem, translate)
-        for (url, problem) in reshape(lines, 2)
-    ]
-    if to_dict:
-      return cls.to_dict(data)
-    return data
-
-  @classmethod
-  def from_txt(cls, data: str, translate: bool = True) -> Problem:
-    """Load a problem from a str object."""
-    url = ''
-    if '\n' in data:
-      url, data = data.split('\n')
-
-    if ' ? ' in data:
-      clauses, goal = data.split(' ? ')
-      goal = Construction.from_txt(goal)
-    else:
-      clauses, goal = data, None
-
-    clauses = clauses.split('; ')
-    problem = Problem(
-        url=url, clauses=[Clause.from_txt(c) for c in clauses], goal=goal
-    )
-    if translate:
-      return problem.translate()
-    return problem
-
-  @classmethod
-  def to_dict(cls, data: list[Problem]) -> dict[str, Problem]:
-    return {p.url: p for p in data}
-
-  def __init__(self, url: str, clauses: list[Clause], goal: Construction):
-    self.url = url
-    self.clauses = clauses
-    self.goal = goal
-
-  def copy(self) -> Problem:
-    return Problem(self.url, list(self.clauses), self.goal)
-
-  def translate(self) -> Problem:  # to single-char point names
-    """Translate point names into alphabetical."""
-    mapping = {}
-    clauses = []
-
-    for clause in self.clauses:
-      clauses.append(clause.translate(mapping))
-
-    if self.goal:
-      goal = self.goal.translate(mapping)
-    else:
-      goal = self.goal
-
-    p = Problem(self.url, clauses, goal)
-    p.mapping = mapping
-    return p
-
-  def txt(self) -> str:
-    return (
-        '; '.join([c.txt() for c in self.clauses]) + ' ? ' + self.goal.txt()
-        if self.goal
-        else ''
-    )
-
-  def setup_str_from_problem(self, definitions: list[Definition]) -> str:
-    """Construct the <theorem_premises> string from Problem object."""
-    ref = 0
-
-    string = []
-    for clause in self.clauses:
-      group = {}
-      p2deps = defaultdict(list)
-      for c in clause.constructions:
-        cdef = definitions[c.name]
-
-        if len(c.args) != len(cdef.construction.args):
-          assert len(c.args) + len(clause.points) == len(cdef.construction.args)
-          c.args = clause.points + c.args
-
-        mapping = dict(zip(cdef.construction.args, c.args))
-        for points, bs in cdef.basics:
-          points = tuple([mapping[x] for x in points])
-          for p in points:
-            group[p] = points
-
-          for b in bs:
-            args = [mapping[a] for a in b.args]
-            name = b.name
-            if b.name in ['s_angle', 'aconst']:
-              x, y, z, v = args
-              name = 'aconst'
-              v = int(v)
-
-              if v < 0:
-                v = -v
-                x, z = z, x
-
-              m, n = simplify(int(v), 180)
-              args = [y, z, y, x, f'{m}pi/{n}']
-
-            p2deps[points].append(hashed_txt(name, args))
-
-      for k, v in p2deps.items():
-        p2deps[k] = sort_deps(v)
-
-      points = clause.points
-      while points:
-        p = points[0]
-        gr = group[p]
-        points = [x for x in points if x not in gr]
-
-        deps_str = []
-        for dep in p2deps[gr]:
-          ref_str = '{:02}'.format(ref)
-          dep_str = pt.pretty(dep)
-
-          if dep[0] == 'aconst':
-            m, n = map(int, dep[-1].split('pi/'))
-            mn = f'{m}. pi / {n}.'
-            dep_str = ' '.join(dep_str.split()[:-1] + [mn])
-
-          deps_str.append(dep_str + ' ' + ref_str)
-          ref += 1
-
-        string.append(' '.join(gr) + ' : ' + ' '.join(deps_str))
-
-    string = '{S} ' + ' ; '.join([s.strip() for s in string])
-    goal = self.goal
-    string += ' ? ' + pt.pretty([goal.name] + goal.args)
-    return string
-
-
-def parse_rely(s: str) -> dict[str, str]:
-  result = {}
-  if not s:
-    return result
-  s = [x.strip() for x in s.split(',')]
-  for x in s:
-    a, b = x.split(':')
-    a, b = a.strip().split(), b.strip().split()
-    result.update({m: b for m in a})
-  return result
-
-
-class Definition:
-  """Definitions of construction statements."""
-
-  @classmethod
-  def from_txt_file(cls, fname: str, to_dict: bool = False) -> Definition:
-    with open(fname, 'r') as f:
-      lines = f.read()
-    return cls.from_string(lines, to_dict)
-
-  @classmethod
-  def from_string(cls, string: str, to_dict: bool = False) -> Definition:
-    lines = string.split('\n')
-    data = [cls.from_txt('\n'.join(group)) for group in reshape(lines, 6)]
-    if to_dict:
-      return cls.to_dict(data)
-    return data
-
-  @classmethod
-  def to_dict(cls, data: list[Definition]) -> dict[str, Definition]:
-    return {d.construction.name: d for d in data}
-
-  @classmethod
-  def from_txt(cls, data: str) -> Definition:
-    """Load definitions from a str object."""
-    construction, rely, deps, basics, numerics, _ = data.split('\n')
-    basics = [] if not basics else [b.strip() for b in basics.split(';')]
-
-    levels = []
-    for bs in basics:
-      if ':' in bs:
-        points, bs = bs.split(':')
-        points = points.strip().split()
-      else:
-        points = []
-      if bs.strip():
-        bs = [Construction.from_txt(b.strip()) for b in bs.strip().split(',')]
-      else:
-        bs = []
-      levels.append((points, bs))
-
-    numerics = [] if not numerics else numerics.split(', ')
-
-    return Definition(
-        construction=Construction.from_txt(construction),
-        rely=parse_rely(rely),
-        deps=Clause.from_txt(deps),
-        basics=levels,
-        numerics=[Construction.from_txt(c) for c in numerics],
-    )
-
-  def __init__(
-      self,
-      construction: Construction,
-      rely: dict[str, str],
-      deps: Clause,
-      basics: list[tuple[list[str], list[Construction]]],
-      numerics: list[Construction],
-  ):
-    self.construction = construction
-    self.rely = rely
-    self.deps = deps
-    self.basics = basics
-    self.numerics = numerics
-
-    args = set()
-    for num in numerics:
-      args.update(num.args)
-
-    self.points = []
-    self.args = []
-    for p in self.construction.args:
-      if p in args:
-        self.args.append(p)
-      else:
-        self.points.append(p)
-
-
-class Theorem:
-  """Deduction rule."""
-
-  @classmethod
-  def from_txt_file(cls, fname: str, to_dict: bool = False) -> Theorem:
-    with open(fname, 'r') as f:
-      theorems = f.read()
-    return cls.from_string(theorems, to_dict)
-
-  @classmethod
-  def from_string(cls, string: str, to_dict: bool = False) -> Theorem:
-    """Load deduction rule from a str object."""
-    theorems = string.split('\n')
-    theorems = [l for l in theorems if l and not l.startswith('#')]
-    theorems = [cls.from_txt(l) for l in theorems]
-
-    for i, th in enumerate(theorems):
-      th.rule_name = 'r{:02}'.format(i)
-
-    if to_dict:
-      result = {}
-      for t in theorems:
-        if t.name in result:
-          t.name += '_'
-        result[t.rule_name] = t
-
-      return result
-
-    return theorems
-
-  @classmethod
-  def from_txt(cls, data: str) -> Theorem:
-    premises, conclusion = data.split(' => ')
-    premises = premises.split(', ')
-    conclusion = conclusion.split(', ')
-    return Theorem(
-        premise=[Construction.from_txt(p) for p in premises],
-        conclusion=[Construction.from_txt(c) for c in conclusion],
-    )
-
-  def __init__(
-      self, premise: list[Construction], conclusion: list[Construction]
-  ):
-    if len(conclusion) != 1:
-      raise ValueError('Cannot have more than one conclusion')
-    self.name = '_'.join([p.name for p in premise + conclusion])
-    self.premise = premise
-    self.conclusion = conclusion
-    self.is_arg_reduce = False
-
-    assert len(self.conclusion) == 1
-    con = self.conclusion[0]
-
-    if con.name in [
-        'eqratio3',
-        'midp',
-        'contri',
-        'simtri',
-        'contri2',
-        'simtri2',
-        'simtri*',
-        'contri*',
-    ]:
-      return
-
-    prem_args = set(sum([p.args for p in self.premise], []))
-    con_args = set(con.args)
-    if len(prem_args) <= len(con_args):
-      self.is_arg_reduce = True
-
-  def txt(self) -> str:
-    premise_txt = ', '.join([clause.txt() for clause in self.premise])
-    conclusion_txt = ', '.join([clause.txt() for clause in self.conclusion])
-    return f'{premise_txt} => {conclusion_txt}'
-
-  def conclusion_name_args(
-      self, mapping: dict[str, gm.Point]
-  ) -> tuple[str, list[gm.Point]]:
-    mapping = {arg: p for arg, p in mapping.items() if isinstance(arg, str)}
-    c = self.conclusion[0]
-    args = [mapping[a] for a in c.args]
-    return c.name, args
-
-
-def why_eqratio(
-    d1: gm.Direction,
-    d2: gm.Direction,
-    d3: gm.Direction,
-    d4: gm.Direction,
-    level: int,
-) -> list[Dependency]:
-  """Why two ratios are equal, returns a Dependency objects."""
-  all12 = list(gm.all_ratios(d1, d2, level))
-  all34 = list(gm.all_ratios(d3, d4, level))
-
-  min_why = None
-  for ang12, d1s, d2s in all12:
-    for ang34, d3s, d4s in all34:
-      why0 = gm.why_equal(ang12, ang34, level)
-      if why0 is None:
-        continue
-      d1_, d2_ = ang12._l
-      d3_, d4_ = ang34._l
-      why1 = gm.bfs_backtrack(d1, [d1_], d1s)
-      why2 = gm.bfs_backtrack(d2, [d2_], d2s)
-      why3 = gm.bfs_backtrack(d3, [d3_], d3s)
-      why4 = gm.bfs_backtrack(d4, [d4_], d4s)
-      why = why0 + why1 + why2 + why3 + why4
-      if min_why is None or len(why) < len(min_why[0]):
-        min_why = why, ang12, ang34, why0, why1, why2, why3, why4
-
-  if min_why is None:
-    return None
-
-  _, ang12, ang34, why0, why1, why2, why3, why4 = min_why
-  d1_, d2_ = ang12._l
-  d3_, d4_ = ang34._l
-
-  if d1 == d1_ and d2 == d2_ and d3 == d3_ and d4 == d4_:
-    return why0
-
-  (a_, b_), (c_, d_) = d1_._obj.points, d2_._obj.points
-  (e_, f_), (g_, h_) = d3_._obj.points, d4_._obj.points
-  deps = []
-  if why0:
-    dep = Dependency('eqratio', [a_, b_, c_, d_, e_, f_, g_, h_], '', level)
-    dep.why = why0
-    deps.append(dep)
-
-  (a, b), (c, d) = d1._obj.points, d2._obj.points
-  (e, f), (g, h) = d3._obj.points, d4._obj.points
-  for why, (x, y), (x_, y_) in zip(
-      [why1, why2, why3, why4],
-      [(a, b), (c, d), (e, f), (g, h)],
-      [(a_, b_), (c_, d_), (e_, f_), (g_, h_)],
-  ):
-    if why:
-      dep = Dependency('cong', [x, y, x_, y_], '', level)
-      dep.why = why
-      deps.append(dep)
-
-  return deps
-
-
-def why_eqangle(
-    d1: gm.Direction,
-    d2: gm.Direction,
-    d3: gm.Direction,
-    d4: gm.Direction,
-    level: int,
-    verbose: bool = False,
-) -> list[Dependency]:
-  """Why two angles are equal, returns a Dependency objects."""
-  all12 = list(gm.all_angles(d1, d2, level))
-  all34 = list(gm.all_angles(d3, d4, level))
-
-  min_why = None
-  for ang12, d1s, d2s in all12:
-    for ang34, d3s, d4s in all34:
-      why0 = gm.why_equal(ang12, ang34, level)
-      if why0 is None:
-        continue
-      d1_, d2_ = ang12._d
-      d3_, d4_ = ang34._d
-      why1 = gm.bfs_backtrack(d1, [d1_], d1s)
-      why2 = gm.bfs_backtrack(d2, [d2_], d2s)
-      why3 = gm.bfs_backtrack(d3, [d3_], d3s)
-      why4 = gm.bfs_backtrack(d4, [d4_], d4s)
-      why = why0 + why1 + why2 + why3 + why4
-      if min_why is None or len(why) < len(min_why[0]):
-        min_why = why, ang12, ang34, why0, why1, why2, why3, why4
-
-  if min_why is None:
-    return None
-
-  _, ang12, ang34, why0, why1, why2, why3, why4 = min_why
-  why0 = gm.why_equal(ang12, ang34, level)
-  d1_, d2_ = ang12._d
-  d3_, d4_ = ang34._d
-
-  if d1 == d1_ and d2 == d2_ and d3 == d3_ and d4 == d4_:
-    return (d1_, d2_, d3_, d4_), why0
-
-  (a_, b_), (c_, d_) = d1_._obj.points, d2_._obj.points
-  (e_, f_), (g_, h_) = d3_._obj.points, d4_._obj.points
-  deps = []
-  if why0:
-    dep = Dependency('eqangle', [a_, b_, c_, d_, e_, f_, g_, h_], '', None)
-    dep.why = why0
-    deps.append(dep)
-
-  (a, b), (c, d) = d1._obj.points, d2._obj.points
-  (e, f), (g, h) = d3._obj.points, d4._obj.points
-  for why, d_xy, (x, y), d_xy_, (x_, y_) in zip(
-      [why1, why2, why3, why4],
-      [d1, d2, d3, d4],
-      [(a, b), (c, d), (e, f), (g, h)],
-      [d1_, d2_, d3_, d4_],
-      [(a_, b_), (c_, d_), (e_, f_), (g_, h_)],
-  ):
-    xy, xy_ = d_xy._obj, d_xy_._obj
-    if why:
-      if xy == xy_:
-        name = 'collx'
-      else:
-        name = 'para'
-      dep = Dependency(name, [x_, y_, x, y], '', None)
-      dep.why = why
-      deps.append(dep)
-
-  return (d1_, d2_, d3_, d4_), deps
-
-
-CONSTRUCTION_RULE = 'c0'
-
-
-class EmptyDependency:
-  """Empty dependency predicate ready to get filled up."""
-
-  def __init__(self, level: int, rule_name: str):
-    self.level = level
-    self.rule_name = rule_name or ''
-    self.empty = True
-    self.why = []
-    self.trace = None
-
-  def populate(self, name: str, args: list[gm.Point]) -> Dependency:
-    dep = Dependency(name, args, self.rule_name, self.level)
-    dep.trace2 = self.trace
-    dep.why = list(self.why)
-    return dep
-
-  def copy(self) -> EmptyDependency:
-    other = EmptyDependency(self.level, self.rule_name)
-    other.why = list(self.why)
-    return other
-
-  def extend(
-      self,
-      g: Any,
-      name0: str,
-      args0: list[gm.Point],
-      name: str,
-      args: list[gm.Point],
-  ) -> EmptyDependency:
-    """Extend the dependency list by (name, args)."""
-    dep0 = self.populate(name0, args0)
-    deps = EmptyDependency(level=self.level, rule_name=None)
-    dep = Dependency(name, args, None, deps.level)
-    deps.why = [dep0, dep.why_me_or_cache(g, None)]
-    return deps
-
-  def extend_many(
-      self,
-      g: Any,
-      name0: str,
-      args0: list[gm.Point],
-      name_args: list[tuple[str, list[gm.Point]]],
-  ) -> EmptyDependency:
-    """Extend the dependency list by many name_args."""
-    if not name_args:
-      return self
-    dep0 = self.populate(name0, args0)
-    deps = EmptyDependency(level=self.level, rule_name=None)
-    deps.why = [dep0]
-    for name, args in name_args:
-      dep = Dependency(name, args, None, deps.level)
-      deps.why += [dep.why_me_or_cache(g, None)]
-    return deps
-
-
-def maybe_make_equal_pairs(
-    a: gm.Point,
-    b: gm.Point,
-    c: gm.Point,
-    d: gm.Point,
-    m: gm.Point,
-    n: gm.Point,
-    p: gm.Point,
-    q: gm.Point,
-    ab: gm.Line,
-    mn: gm.Line,
-    g: Any,
-    level: int,
-) -> list[Dependency]:
-  """Make a-b:c-d==m-n:p-q in case a-b==m-n or c-d==p-q."""
-  if ab != mn:
-    return
-  why = []
-  eqname = 'para' if isinstance(ab, gm.Line) else 'cong'
-  colls = [a, b, m, n]
-  if len(set(colls)) > 2 and eqname == 'para':
-    dep = Dependency('collx', colls, None, level)
-    dep.why_me(g, level)
-    why += [dep]
-
-  dep = Dependency(eqname, [c, d, p, q], None, level)
-  dep.why_me(g, level)
-  why += [dep]
-  return why
-
-
-class Dependency(Construction):
-  """Dependency is a predicate that other predicates depend on."""
-
-  def __init__(
-      self, name: str, args: list[gm.Point], rule_name: str, level: int
-  ):
-    super().__init__(name, args)
-    self.rule_name = rule_name or ''
-    self.level = level
-    self.why = []
-
-    self._stat = None
-    self.trace = None
-
-  def _find(self, dep_hashed: tuple[str, ...]) -> Dependency:
-    for w in self.why:
-      f = w._find(dep_hashed)
-      if f:
-        return f
-      if w.hashed() == dep_hashed:
-        return w
-
-  def remove_loop(self) -> Dependency:
-    f = self._find(self.hashed())
-    if f:
-      return f
-    return self
-
-  def copy(self) -> Dependency:
-    dep = Dependency(self.name, self.args, self.rule_name, self.level)
-    dep.trace = self.trace
-    dep.why = list(self.why)
-    return dep
-
-  def why_me_or_cache(self, g: Any, level: int) -> Dependency:
-    if self.hashed() in g.cache:
-      return g.cache[self.hashed()]
-    self.why_me(g, level)
-    return self
-
-  def populate(self, name: str, args: list[gm.Point]) -> Dependency:
-    assert self.rule_name == CONSTRUCTION_RULE, self.rule_name
-    dep = Dependency(self.name, self.args, self.rule_name, self.level)
-    dep.why = list(self.why)
-    return dep
-
-  def why_me(self, g: Any, level: int) -> None:
-    """Figure out the dependencies predicates of self."""
-    name, args = self.name, self.args
-
-    hashed_me = hashed(name, args)
-    if hashed_me in g.cache:
-      dep = g.cache[hashed_me]
-      self.why = dep.why
-      self.rule_name = dep.rule_name
-      return
-
-    if self.name == 'para':
-      a, b, c, d = self.args
-      if {a, b} == {c, d}:
-        self.why = []
-        return
-
-      ab = g._get_line(a, b)
-      cd = g._get_line(c, d)
-      if ab == cd:
-        if {a, b} == {c, d}:
-          self.why = []
-          self.rule_name = ''
-          return
-        dep = Dependency('coll', list({a, b, c, d}), 't??', None)
-        self.why = [dep.why_me_or_cache(g, level)]
-        return
-
-      for (x, y), xy in zip([(a, b), (c, d)], [ab, cd]):
-        x_, y_ = xy.points
-        if {x, y} == {x_, y_}:
-          continue
-        d = Dependency('collx', [x, y, x_, y_], None, level)
-        self.why += [d.why_me_or_cache(g, level)]
-
-      whypara = g.why_equal(ab, cd, None)
-      self.why += whypara
-
-    elif self.name == 'midp':
-      m, a, b = self.args
-      ma = g._get_segment(m, a)
-      mb = g._get_segment(m, b)
-      dep = Dependency('coll', [m, a, b], None, None).why_me_or_cache(g, None)
-      self.why = [dep] + g.why_equal(ma, mb, level)
-
-    elif self.name == 'perp':
-      a, b, c, d = self.args
-      ab = g._get_line(a, b)
-      cd = g._get_line(c, d)
-      for (x, y), xy in zip([(a, b), (c, d)], [ab, cd]):
-        x_, y_ = xy.points
-        if {x, y} == {x_, y_}:
-          continue
-        d = Dependency('collx', [x, y, x_, y_], None, level)
-        self.why += [d.why_me_or_cache(g, level)]
-
-      _, why = why_eqangle(ab._val, cd._val, cd._val, ab._val, level)
-      a, b = ab.points
-      c, d = cd.points
-
-      if hashed(self.name, [a, b, c, d]) != self.hashed():
-        d = Dependency(self.name, [a, b, c, d], None, level)
-        d.why = why
-        why = [d]
-
-      self.why += why
-
-    elif self.name == 'cong':
-      a, b, c, d = self.args
-      ab = g._get_segment(a, b)
-      cd = g._get_segment(c, d)
-
-      self.why = g.why_equal(ab, cd, level)
-
-    elif self.name == 'coll':
-      _, why = gm.line_of_and_why(self.args, level)
-      self.why = why
-
-    elif self.name == 'collx':
-      if g.check_coll(self.args):
-        args = list(set(self.args))
-        hashed_me = hashed('coll', args)
-        if hashed_me in g.cache:
-          dep = g.cache[hashed_me]
-          self.why = [dep]
-          self.rule_name = ''
-          return
-        _, self.why = gm.line_of_and_why(args, level)
-      else:
-        self.name = 'para'
-        self.why_me(g, level)
-
-    elif self.name == 'cyclic':
-      _, why = gm.circle_of_and_why(self.args, level)
-      self.why = why
-
-    elif self.name == 'circle':
-      o, a, b, c = self.args
-      oa = g._get_segment(o, a)
-      ob = g._get_segment(o, b)
-      oc = g._get_segment(o, c)
-      self.why = g.why_equal(oa, ob, level) + g.why_equal(oa, oc, level)
-
-    elif self.name in ['eqangle', 'eqangle6']:
-      a, b, c, d, m, n, p, q = self.args
-
-      ab, why1 = g.get_line_thru_pair_why(a, b)
-      cd, why2 = g.get_line_thru_pair_why(c, d)
-      mn, why3 = g.get_line_thru_pair_why(m, n)
-      pq, why4 = g.get_line_thru_pair_why(p, q)
-
-      if ab is None or cd is None or mn is None or pq is None:
-        if {a, b} == {m, n}:
-          d = Dependency('para', [c, d, p, q], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {a, b} == {c, d}:
-          d = Dependency('para', [p, q, m, n], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {c, d} == {p, q}:
-          d = Dependency('para', [a, b, m, n], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {p, q} == {m, n}:
-          d = Dependency('para', [a, b, c, d], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        return
-
-      for (x, y), xy, whyxy in zip(
-          [(a, b), (c, d), (m, n), (p, q)],
-          [ab, cd, mn, pq],
-          [why1, why2, why3, why4],
-      ):
-        x_, y_ = xy.points
-        if {x, y} == {x_, y_}:
-          continue
-        d = Dependency('collx', [x, y, x_, y_], None, level)
-        d.why = whyxy
-        self.why += [d]
-
-      a, b = ab.points
-      c, d = cd.points
-      m, n = mn.points
-      p, q = pq.points
-      diff = hashed(self.name, [a, b, c, d, m, n, p, q]) != self.hashed()
-
-      whyeqangle = None
-      if ab._val and cd._val and mn._val and pq._val:
-        whyeqangle = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
-
-      if whyeqangle:
-        (dab, dcd, dmn, dpq), whyeqangle = whyeqangle
-        if diff:
-          d = Dependency('eqangle', [a, b, c, d, m, n, p, q], None, level)
-          d.why = whyeqangle
-          whyeqangle = [d]
-        self.why += whyeqangle
-
-      else:
-        if (ab == cd and mn == pq) or (ab == mn and cd == pq):
-          self.why += []
-        elif ab == mn:
-          self.why += maybe_make_equal_pairs(
-              a, b, c, d, m, n, p, q, ab, mn, g, level
-          )
-        elif cd == pq:
-          self.why += maybe_make_equal_pairs(
-              c, d, a, b, p, q, m, n, cd, pq, g, level
-          )
-        elif ab == cd:
-          self.why += maybe_make_equal_pairs(
-              a, b, m, n, c, d, p, q, ab, cd, g, level
-          )
-        elif mn == pq:
-          self.why += maybe_make_equal_pairs(
-              m, n, a, b, p, q, c, d, mn, pq, g, level
-          )
-        elif g.is_equal(ab, mn) or g.is_equal(cd, pq):
-          dep1 = Dependency('para', [a, b, m, n], None, level)
-          dep1.why_me(g, level)
-          dep2 = Dependency('para', [c, d, p, q], None, level)
-          dep2.why_me(g, level)
-          self.why += [dep1, dep2]
-        elif g.is_equal(ab, cd) or g.is_equal(mn, pq):
-          dep1 = Dependency('para', [a, b, c, d], None, level)
-          dep1.why_me(g, level)
-          dep2 = Dependency('para', [m, n, p, q], None, level)
-          dep2.why_me(g, level)
-          self.why += [dep1, dep2]
-        elif ab._val and cd._val and mn._val and pq._val:
-          self.why = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
-
-    elif self.name in ['eqratio', 'eqratio6']:
-      a, b, c, d, m, n, p, q = self.args
-      ab = g._get_segment(a, b)
-      cd = g._get_segment(c, d)
-      mn = g._get_segment(m, n)
-      pq = g._get_segment(p, q)
-
-      if ab is None or cd is None or mn is None or pq is None:
-        if {a, b} == {m, n}:
-          d = Dependency('cong', [c, d, p, q], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {a, b} == {c, d}:
-          d = Dependency('cong', [p, q, m, n], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {c, d} == {p, q}:
-          d = Dependency('cong', [a, b, m, n], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        if {p, q} == {m, n}:
-          d = Dependency('cong', [a, b, c, d], None, level)
-          self.why = [d.why_me_or_cache(g, level)]
-        return
-
-      if ab._val and cd._val and mn._val and pq._val:
-        self.why = why_eqratio(ab._val, cd._val, mn._val, pq._val, level)
-
-      if self.why is None:
-        self.why = []
-        if (ab == cd and mn == pq) or (ab == mn and cd == pq):
-          self.why = []
-        elif ab == mn:
-          self.why += maybe_make_equal_pairs(
-              a, b, c, d, m, n, p, q, ab, mn, g, level
-          )
-        elif cd == pq:
-          self.why += maybe_make_equal_pairs(
-              c, d, a, b, p, q, m, n, cd, pq, g, level
-          )
-        elif ab == cd:
-          self.why += maybe_make_equal_pairs(
-              a, b, m, n, c, d, p, q, ab, cd, g, level
-          )
-        elif mn == pq:
-          self.why += maybe_make_equal_pairs(
-              m, n, a, b, p, q, c, d, mn, pq, g, level
-          )
-        elif g.is_equal(ab, mn) or g.is_equal(cd, pq):
-          dep1 = Dependency('cong', [a, b, m, n], None, level)
-          dep1.why_me(g, level)
-          dep2 = Dependency('cong', [c, d, p, q], None, level)
-          dep2.why_me(g, level)
-          self.why += [dep1, dep2]
-        elif g.is_equal(ab, cd) or g.is_equal(mn, pq):
-          dep1 = Dependency('cong', [a, b, c, d], None, level)
-          dep1.why_me(g, level)
-          dep2 = Dependency('cong', [m, n, p, q], None, level)
-          dep2.why_me(g, level)
-          self.why += [dep1, dep2]
-        elif ab._val and cd._val and mn._val and pq._val:
-          self.why = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
-
-    elif self.name in ['diff', 'npara', 'nperp', 'ncoll', 'sameside']:
-      self.why = []
-
-    elif self.name == 'simtri':
-      a, b, c, x, y, z = self.args
-      dep1 = Dependency('eqangle', [a, b, a, c, x, y, x, z], '', level)
-      dep1.why_me(g, level)
-      dep2 = Dependency('eqangle', [b, a, b, c, y, x, y, z], '', level)
-      dep2.why_me(g, level)
-      self.rule_name = 'r34'
-      self.why = [dep1, dep2]
-
-    elif self.name == 'contri':
-      a, b, c, x, y, z = self.args
-      dep1 = Dependency('cong', [a, b, x, y], '', level)
-      dep1.why_me(g, level)
-      dep2 = Dependency('cong', [b, c, y, z], '', level)
-      dep2.why_me(g, level)
-      dep3 = Dependency('cong', [c, a, z, x], '', level)
-      dep3.why_me(g, level)
-      self.rule_name = 'r32'
-      self.why = [dep1, dep2, dep3]
-
-    elif self.name == 'ind':
-      pass
-
-    elif self.name == 'aconst':
-      a, b, c, d, ang0 = self.args
-
-      measure = ang0._val
-
-      for ang in measure.neighbors(gm.Angle):
-        if ang == ang0:
-          continue
-        d1, d2 = ang._d
-        l1, l2 = d1._obj, d2._obj
-        (a1, b1), (c1, d1) = l1.points, l2.points
-
-        if not g.check_para_or_coll([a, b, a1, b1]) or not g.check_para_or_coll(
-            [c, d, c1, d1]
-        ):
-          continue
-
-        self.why = []
-        for args in [(a, b, a1, b1), (c, d, c1, d1)]:
-          if g.check_coll(args):
-            if len(set(args)) > 2:
-              dep = Dependency('coll', args, None, None)
-              self.why.append(dep.why_me_or_cache(g, level))
-          else:
-            dep = Dependency('para', args, None, None)
-            self.why.append(dep.why_me_or_cache(g, level))
-
-        self.why += gm.why_equal(ang, ang0)
-        break
-
-    elif self.name == 'rconst':
-      a, b, c, d, rat0 = self.args
-
-      val = rat0._val
-
-      for rat in val.neighbors(gm.Ratio):
-        if rat == rat0:
-          continue
-        l1, l2 = rat._l
-        s1, s2 = l1._obj, l2._obj
-        (a1, b1), (c1, d1) = list(s1.points), list(s2.points)
-
-        if not g.check_cong([a, b, a1, b1]) or not g.check_cong([c, d, c1, d1]):
-          continue
-
-        self.why = []
-        for args in [(a, b, a1, b1), (c, d, c1, d1)]:
-          if len(set(args)) > 2:
-            dep = Dependency('cong', args, None, None)
-            self.why.append(dep.why_me_or_cache(g, level))
-
-        self.why += gm.why_equal(rat, rat0)
-        break
-
-    else:
-      raise ValueError('Not recognize', self.name)
-
-  def hashed(self, rename: bool = False) -> tuple[str, ...]:
-    return hashed(self.name, self.args, rename=rename)
-
-
-def hashed(
-    name: str, args: list[gm.Point], rename: bool = False
-) -> tuple[str, ...]:
-  if name == 's_angle':
-    args = [p.name if not rename else p.new_name for p in args[:-1]] + [
-        str(args[-1])
-    ]
-  else:
-    args = [p.name if not rename else p.new_name for p in args]
-  return hashed_txt(name, args)
-
-
-def hashed_txt(name: str, args: list[str]) -> tuple[str, ...]:
-  """Return a tuple unique to name and args upto arg permutation equivariant."""
-
-  if name in ['const', 'aconst', 'rconst']:
-    a, b, c, d, y = args
-    a, b = sorted([a, b])
-    c, d = sorted([c, d])
-    return name, a, b, c, d, y
-
-  if name in ['npara', 'nperp', 'para', 'cong', 'perp', 'collx']:
-    a, b, c, d = args
-
-    a, b = sorted([a, b])
-    c, d = sorted([c, d])
-    (a, b), (c, d) = sorted([(a, b), (c, d)])
-
-    return (name, a, b, c, d)
-
-  if name in ['midp', 'midpoint']:
-    a, b, c = args
-    b, c = sorted([b, c])
-    return (name, a, b, c)
-
-  if name in ['coll', 'cyclic', 'ncoll', 'diff', 'triangle']:
-    return (name,) + tuple(sorted(list(set(args))))
-
-  if name == 'circle':
-    x, a, b, c = args
-    return (name, x) + tuple(sorted([a, b, c]))
-
-  if name in ['eqangle', 'eqratio', 'eqangle6', 'eqratio6']:
-    a, b, c, d, e, f, g, h = args
-    a, b = sorted([a, b])
-    c, d = sorted([c, d])
-    e, f = sorted([e, f])
-    g, h = sorted([g, h])
-    if tuple(sorted([a, b, e, f])) > tuple(sorted([c, d, g, h])):
-      a, b, e, f, c, d, g, h = c, d, g, h, a, b, e, f
-    if (a, b, c, d) > (e, f, g, h):
-      a, b, c, d, e, f, g, h = e, f, g, h, a, b, c, d
-
-    if name == 'eqangle6':
-      name = 'eqangle'
-    if name == 'eqratio6':
-      name = 'eqratio'
-    return (name,) + (a, b, c, d, e, f, g, h)
-
-  if name in ['contri', 'simtri', 'simtri2', 'contri2', 'contri*', 'simtri*']:
-    a, b, c, x, y, z = args
-    (a, x), (b, y), (c, z) = sorted([(a, x), (b, y), (c, z)], key=sorted)
-    (a, b, c), (x, y, z) = sorted([(a, b, c), (x, y, z)], key=sorted)
-    return (name, a, b, c, x, y, z)
-
-  if name in ['eqratio3']:
-    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
-    (a, c), (b, d) = sorted([(a, c), (b, d)], key=sorted)
-    (a, b), (c, d) = sorted([(a, b), (c, d)], key=sorted)
-    return (name, a, b, c, d, o, o)
-
-  if name in ['sameside', 's_angle']:
-    return (name,) + tuple(args)
-
-  raise ValueError(f'Not recognize {name} to hash.')
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Implements objects to represent problems, theorems, proofs, traceback."""
+
+from __future__ import annotations
+
+from collections import defaultdict  # pylint: disable=g-importing-member
+from typing import Any
+
+import geometry as gm
+import pretty as pt
+
+
+# pylint: disable=protected-access
+# pylint: disable=unused-variable
+# pylint: disable=unused-argument
+# pylint: disable=unused-assignment
+
+
+def reshape(l: list[Any], n: int = 1) -> list[list[Any]]:
+  assert len(l) % n == 0
+  columns = [[] for i in range(n)]
+  for i, x in enumerate(l):
+    columns[i % n].append(x)
+  return zip(*columns)
+
+
+def isint(x: str) -> bool:
+  try:
+    int(x)
+    return True
+  except:  # pylint: disable=bare-except
+    return False
+
+
+class Construction:
+  """One predicate."""
+
+  @classmethod
+  def from_txt(cls, data: str) -> Construction:
+    data = data.split(' ')
+    return Construction(data[0], data[1:])
+
+  def __init__(self, name: str, args: list[str]):
+    self.name = name
+    self.args = args
+
+  def translate(self, mapping: dict[str, str]) -> Construction:
+    args = [a if isint(a) else mapping[a] for a in self.args]
+    return Construction(self.name, args)
+
+  def txt(self) -> str:
+    return ' '.join([self.name] + list(self.args))
+
+
+class Clause:
+  """One construction (>= 1 predicate)."""
+
+  @classmethod
+  def from_txt(cls, data: str) -> Clause:
+    if data == ' =':
+      return Clause([], [])
+    points, constructions = data.split(' = ')
+    return Clause(
+        points.split(' '),
+        [Construction.from_txt(c) for c in constructions.split(', ')],
+    )
+
+  def __init__(self, points: list[str], constructions: list[Construction]):
+    self.points = []
+    self.nums = []
+
+    for p in points:
+      num = None
+      if isinstance(p, str) and '@' in p:
+        p, num = p.split('@')
+        x, y = num.split('_')
+        num = float(x), float(y)
+      self.points.append(p)
+      self.nums.append(num)
+
+    self.constructions = constructions
+
+  def translate(self, mapping: dict[str, str]) -> Clause:
+    points0 = []
+    for p in self.points:
+      pcount = len(mapping) + 1
+      name = chr(96 + pcount)
+      if name > 'z':  # pcount = 26 -> name = 'z'
+        name = chr(97 + (pcount - 1) % 26) + str((pcount - 1) // 26)
+
+      p0 = mapping.get(p, name)
+      mapping[p] = p0
+      points0.append(p0)
+    return Clause(points0, [c.translate(mapping) for c in self.constructions])
+
+  def add(self, name: str, args: list[str]) -> None:
+    self.constructions.append(Construction(name, args))
+
+  def txt(self) -> str:
+    return (
+        ' '.join(self.points)
+        + ' = '
+        + ', '.join(c.txt() for c in self.constructions)
+    )
+
+
+def _gcd(x: int, y: int) -> int:
+  while y:
+    x, y = y, x % y
+  return x
+
+
+def simplify(n: int, d: int) -> tuple[int, int]:
+  g = _gcd(n, d)
+  return (n // g, d // g)
+
+
+def compare_fn(dep: Dependency) -> tuple[Dependency, str]:
+  return (dep, pt.pretty(dep))
+
+
+def sort_deps(deps: list[Dependency]) -> list[Dependency]:
+  return sorted(deps, key=compare_fn)
+
+
+class Problem:
+  """Describe one problem to solve."""
+
+  @classmethod
+  def from_txt_file(
+      cls, fname: str, to_dict: bool = False, translate: bool = True
+  ):
+    """Load a problem from a text file."""
+    with open(fname, 'r') as f:
+      lines = f.read().split('\n')
+
+    lines = [l for l in lines if l]
+    data = [
+        cls.from_txt(url + '\n' + problem, translate)
+        for (url, problem) in reshape(lines, 2)
+    ]
+    if to_dict:
+      return cls.to_dict(data)
+    return data
+
+  @classmethod
+  def from_txt(cls, data: str, translate: bool = True) -> Problem:
+    """Load a problem from a str object."""
+    url = ''
+    if '\n' in data:
+      url, data = data.split('\n')
+
+    if ' ? ' in data:
+      clauses, goal = data.split(' ? ')
+      goal = Construction.from_txt(goal)
+    else:
+      clauses, goal = data, None
+
+    clauses = clauses.split('; ')
+    problem = Problem(
+        url=url, clauses=[Clause.from_txt(c) for c in clauses], goal=goal
+    )
+    if translate:
+      return problem.translate()
+    return problem
+
+  @classmethod
+  def to_dict(cls, data: list[Problem]) -> dict[str, Problem]:
+    return {p.url: p for p in data}
+
+  def __init__(self, url: str, clauses: list[Clause], goal: Construction):
+    self.url = url
+    self.clauses = clauses
+    self.goal = goal
+
+  def copy(self) -> Problem:
+    return Problem(self.url, list(self.clauses), self.goal)
+
+  def translate(self) -> Problem:  # to single-char point names
+    """Translate point names into alphabetical."""
+    mapping = {}
+    clauses = []
+
+    for clause in self.clauses:
+      clauses.append(clause.translate(mapping))
+
+    if self.goal:
+      goal = self.goal.translate(mapping)
+    else:
+      goal = self.goal
+
+    p = Problem(self.url, clauses, goal)
+    p.mapping = mapping
+    return p
+
+  def txt(self) -> str:
+    return (
+        '; '.join([c.txt() for c in self.clauses]) + ' ? ' + self.goal.txt()
+        if self.goal
+        else ''
+    )
+
+  def setup_str_from_problem(self, definitions: list[Definition]) -> str:
+    """Construct the <theorem_premises> string from Problem object."""
+    ref = 0
+
+    string = []
+    for clause in self.clauses:
+      group = {}
+      p2deps = defaultdict(list)
+      for c in clause.constructions:
+        cdef = definitions[c.name]
+
+        if len(c.args) != len(cdef.construction.args):
+          assert len(c.args) + len(clause.points) == len(cdef.construction.args)
+          c.args = clause.points + c.args
+
+        mapping = dict(zip(cdef.construction.args, c.args))
+        for points, bs in cdef.basics:
+          points = tuple([mapping[x] for x in points])
+          for p in points:
+            group[p] = points
+
+          for b in bs:
+            args = [mapping[a] for a in b.args]
+            name = b.name
+            if b.name in ['s_angle', 'aconst']:
+              x, y, z, v = args
+              name = 'aconst'
+              v = int(v)
+
+              if v < 0:
+                v = -v
+                x, z = z, x
+
+              m, n = simplify(int(v), 180)
+              args = [y, z, y, x, f'{m}pi/{n}']
+
+            p2deps[points].append(hashed_txt(name, args))
+
+      for k, v in p2deps.items():
+        p2deps[k] = sort_deps(v)
+
+      points = clause.points
+      while points:
+        p = points[0]
+        gr = group[p]
+        points = [x for x in points if x not in gr]
+
+        deps_str = []
+        for dep in p2deps[gr]:
+          ref_str = '{:02}'.format(ref)
+          dep_str = pt.pretty(dep)
+
+          if dep[0] == 'aconst':
+            m, n = map(int, dep[-1].split('pi/'))
+            mn = f'{m}. pi / {n}.'
+            dep_str = ' '.join(dep_str.split()[:-1] + [mn])
+
+          deps_str.append(dep_str + ' ' + ref_str)
+          ref += 1
+
+        string.append(' '.join(gr) + ' : ' + ' '.join(deps_str))
+
+    string = '{S} ' + ' ; '.join([s.strip() for s in string])
+    goal = self.goal
+    string += ' ? ' + pt.pretty([goal.name] + goal.args)
+    return string
+
+
+def parse_rely(s: str) -> dict[str, str]:
+  result = {}
+  if not s:
+    return result
+  s = [x.strip() for x in s.split(',')]
+  for x in s:
+    a, b = x.split(':')
+    a, b = a.strip().split(), b.strip().split()
+    result.update({m: b for m in a})
+  return result
+
+
+class Definition:
+  """Definitions of construction statements."""
+
+  @classmethod
+  def from_txt_file(cls, fname: str, to_dict: bool = False) -> Definition:
+    with open(fname, 'r') as f:
+      lines = f.read()
+    return cls.from_string(lines, to_dict)
+
+  @classmethod
+  def from_string(cls, string: str, to_dict: bool = False) -> Definition:
+    lines = string.split('\n')
+    data = [cls.from_txt('\n'.join(group)) for group in reshape(lines, 6)]
+    if to_dict:
+      return cls.to_dict(data)
+    return data
+
+  @classmethod
+  def to_dict(cls, data: list[Definition]) -> dict[str, Definition]:
+    return {d.construction.name: d for d in data}
+
+  @classmethod
+  def from_txt(cls, data: str) -> Definition:
+    """Load definitions from a str object."""
+    construction, rely, deps, basics, numerics, _ = data.split('\n')
+    basics = [] if not basics else [b.strip() for b in basics.split(';')]
+
+    levels = []
+    for bs in basics:
+      if ':' in bs:
+        points, bs = bs.split(':')
+        points = points.strip().split()
+      else:
+        points = []
+      if bs.strip():
+        bs = [Construction.from_txt(b.strip()) for b in bs.strip().split(',')]
+      else:
+        bs = []
+      levels.append((points, bs))
+
+    numerics = [] if not numerics else numerics.split(', ')
+
+    return Definition(
+        construction=Construction.from_txt(construction),
+        rely=parse_rely(rely),
+        deps=Clause.from_txt(deps),
+        basics=levels,
+        numerics=[Construction.from_txt(c) for c in numerics],
+    )
+
+  def __init__(
+      self,
+      construction: Construction,
+      rely: dict[str, str],
+      deps: Clause,
+      basics: list[tuple[list[str], list[Construction]]],
+      numerics: list[Construction],
+  ):
+    self.construction = construction
+    self.rely = rely
+    self.deps = deps
+    self.basics = basics
+    self.numerics = numerics
+
+    args = set()
+    for num in numerics:
+      args.update(num.args)
+
+    self.points = []
+    self.args = []
+    for p in self.construction.args:
+      if p in args:
+        self.args.append(p)
+      else:
+        self.points.append(p)
+
+
+class Theorem:
+  """Deduction rule."""
+
+  @classmethod
+  def from_txt_file(cls, fname: str, to_dict: bool = False) -> Theorem:
+    with open(fname, 'r') as f:
+      theorems = f.read()
+    return cls.from_string(theorems, to_dict)
+
+  @classmethod
+  def from_string(cls, string: str, to_dict: bool = False) -> Theorem:
+    """Load deduction rule from a str object."""
+    theorems = string.split('\n')
+    theorems = [l for l in theorems if l and not l.startswith('#')]
+    theorems = [cls.from_txt(l) for l in theorems]
+
+    for i, th in enumerate(theorems):
+      th.rule_name = 'r{:02}'.format(i)
+
+    if to_dict:
+      result = {}
+      for t in theorems:
+        if t.name in result:
+          t.name += '_'
+        result[t.rule_name] = t
+
+      return result
+
+    return theorems
+
+  @classmethod
+  def from_txt(cls, data: str) -> Theorem:
+    premises, conclusion = data.split(' => ')
+    premises = premises.split(', ')
+    conclusion = conclusion.split(', ')
+    return Theorem(
+        premise=[Construction.from_txt(p) for p in premises],
+        conclusion=[Construction.from_txt(c) for c in conclusion],
+    )
+
+  def __init__(
+      self, premise: list[Construction], conclusion: list[Construction]
+  ):
+    if len(conclusion) != 1:
+      raise ValueError('Cannot have more than one conclusion')
+    self.name = '_'.join([p.name for p in premise + conclusion])
+    self.premise = premise
+    self.conclusion = conclusion
+    self.is_arg_reduce = False
+
+    assert len(self.conclusion) == 1
+    con = self.conclusion[0]
+
+    if con.name in [
+        'eqratio3',
+        'midp',
+        'contri',
+        'simtri',
+        'contri2',
+        'simtri2',
+        'simtri*',
+        'contri*',
+    ]:
+      return
+
+    prem_args = set(sum([p.args for p in self.premise], []))
+    con_args = set(con.args)
+    if len(prem_args) <= len(con_args):
+      self.is_arg_reduce = True
+
+  def txt(self) -> str:
+    premise_txt = ', '.join([clause.txt() for clause in self.premise])
+    conclusion_txt = ', '.join([clause.txt() for clause in self.conclusion])
+    return f'{premise_txt} => {conclusion_txt}'
+
+  def conclusion_name_args(
+      self, mapping: dict[str, gm.Point]
+  ) -> tuple[str, list[gm.Point]]:
+    mapping = {arg: p for arg, p in mapping.items() if isinstance(arg, str)}
+    c = self.conclusion[0]
+    args = [mapping[a] for a in c.args]
+    return c.name, args
+
+
+def why_eqratio(
+    d1: gm.Direction,
+    d2: gm.Direction,
+    d3: gm.Direction,
+    d4: gm.Direction,
+    level: int,
+) -> list[Dependency]:
+  """Why two ratios are equal, returns a Dependency objects."""
+  all12 = list(gm.all_ratios(d1, d2, level))
+  all34 = list(gm.all_ratios(d3, d4, level))
+
+  min_why = None
+  for ang12, d1s, d2s in all12:
+    for ang34, d3s, d4s in all34:
+      why0 = gm.why_equal(ang12, ang34, level)
+      if why0 is None:
+        continue
+      d1_, d2_ = ang12._l
+      d3_, d4_ = ang34._l
+      why1 = gm.bfs_backtrack(d1, [d1_], d1s)
+      why2 = gm.bfs_backtrack(d2, [d2_], d2s)
+      why3 = gm.bfs_backtrack(d3, [d3_], d3s)
+      why4 = gm.bfs_backtrack(d4, [d4_], d4s)
+      why = why0 + why1 + why2 + why3 + why4
+      if min_why is None or len(why) < len(min_why[0]):
+        min_why = why, ang12, ang34, why0, why1, why2, why3, why4
+
+  if min_why is None:
+    return None
+
+  _, ang12, ang34, why0, why1, why2, why3, why4 = min_why
+  d1_, d2_ = ang12._l
+  d3_, d4_ = ang34._l
+
+  if d1 == d1_ and d2 == d2_ and d3 == d3_ and d4 == d4_:
+    return why0
+
+  (a_, b_), (c_, d_) = d1_._obj.points, d2_._obj.points
+  (e_, f_), (g_, h_) = d3_._obj.points, d4_._obj.points
+  deps = []
+  if why0:
+    dep = Dependency('eqratio', [a_, b_, c_, d_, e_, f_, g_, h_], '', level)
+    dep.why = why0
+    deps.append(dep)
+
+  (a, b), (c, d) = d1._obj.points, d2._obj.points
+  (e, f), (g, h) = d3._obj.points, d4._obj.points
+  for why, (x, y), (x_, y_) in zip(
+      [why1, why2, why3, why4],
+      [(a, b), (c, d), (e, f), (g, h)],
+      [(a_, b_), (c_, d_), (e_, f_), (g_, h_)],
+  ):
+    if why:
+      dep = Dependency('cong', [x, y, x_, y_], '', level)
+      dep.why = why
+      deps.append(dep)
+
+  return deps
+
+
+def why_eqangle(
+    d1: gm.Direction,
+    d2: gm.Direction,
+    d3: gm.Direction,
+    d4: gm.Direction,
+    level: int,
+    verbose: bool = False,
+) -> list[Dependency]:
+  """Why two angles are equal, returns a Dependency objects."""
+  all12 = list(gm.all_angles(d1, d2, level))
+  all34 = list(gm.all_angles(d3, d4, level))
+
+  min_why = None
+  for ang12, d1s, d2s in all12:
+    for ang34, d3s, d4s in all34:
+      why0 = gm.why_equal(ang12, ang34, level)
+      if why0 is None:
+        continue
+      d1_, d2_ = ang12._d
+      d3_, d4_ = ang34._d
+      why1 = gm.bfs_backtrack(d1, [d1_], d1s)
+      why2 = gm.bfs_backtrack(d2, [d2_], d2s)
+      why3 = gm.bfs_backtrack(d3, [d3_], d3s)
+      why4 = gm.bfs_backtrack(d4, [d4_], d4s)
+      why = why0 + why1 + why2 + why3 + why4
+      if min_why is None or len(why) < len(min_why[0]):
+        min_why = why, ang12, ang34, why0, why1, why2, why3, why4
+
+  if min_why is None:
+    return None
+
+  _, ang12, ang34, why0, why1, why2, why3, why4 = min_why
+  why0 = gm.why_equal(ang12, ang34, level)
+  d1_, d2_ = ang12._d
+  d3_, d4_ = ang34._d
+
+  if d1 == d1_ and d2 == d2_ and d3 == d3_ and d4 == d4_:
+    return (d1_, d2_, d3_, d4_), why0
+
+  (a_, b_), (c_, d_) = d1_._obj.points, d2_._obj.points
+  (e_, f_), (g_, h_) = d3_._obj.points, d4_._obj.points
+  deps = []
+  if why0:
+    dep = Dependency('eqangle', [a_, b_, c_, d_, e_, f_, g_, h_], '', None)
+    dep.why = why0
+    deps.append(dep)
+
+  (a, b), (c, d) = d1._obj.points, d2._obj.points
+  (e, f), (g, h) = d3._obj.points, d4._obj.points
+  for why, d_xy, (x, y), d_xy_, (x_, y_) in zip(
+      [why1, why2, why3, why4],
+      [d1, d2, d3, d4],
+      [(a, b), (c, d), (e, f), (g, h)],
+      [d1_, d2_, d3_, d4_],
+      [(a_, b_), (c_, d_), (e_, f_), (g_, h_)],
+  ):
+    xy, xy_ = d_xy._obj, d_xy_._obj
+    if why:
+      if xy == xy_:
+        name = 'collx'
+      else:
+        name = 'para'
+      dep = Dependency(name, [x_, y_, x, y], '', None)
+      dep.why = why
+      deps.append(dep)
+
+  return (d1_, d2_, d3_, d4_), deps
+
+
+CONSTRUCTION_RULE = 'c0'
+
+
+class EmptyDependency:
+  """Empty dependency predicate ready to get filled up."""
+
+  def __init__(self, level: int, rule_name: str):
+    self.level = level
+    self.rule_name = rule_name or ''
+    self.empty = True
+    self.why = []
+    self.trace = None
+
+  def populate(self, name: str, args: list[gm.Point]) -> Dependency:
+    dep = Dependency(name, args, self.rule_name, self.level)
+    dep.trace2 = self.trace
+    dep.why = list(self.why)
+    return dep
+
+  def copy(self) -> EmptyDependency:
+    other = EmptyDependency(self.level, self.rule_name)
+    other.why = list(self.why)
+    return other
+
+  def extend(
+      self,
+      g: Any,
+      name0: str,
+      args0: list[gm.Point],
+      name: str,
+      args: list[gm.Point],
+  ) -> EmptyDependency:
+    """Extend the dependency list by (name, args)."""
+    dep0 = self.populate(name0, args0)
+    deps = EmptyDependency(level=self.level, rule_name=None)
+    dep = Dependency(name, args, None, deps.level)
+    deps.why = [dep0, dep.why_me_or_cache(g, None)]
+    return deps
+
+  def extend_many(
+      self,
+      g: Any,
+      name0: str,
+      args0: list[gm.Point],
+      name_args: list[tuple[str, list[gm.Point]]],
+  ) -> EmptyDependency:
+    """Extend the dependency list by many name_args."""
+    if not name_args:
+      return self
+    dep0 = self.populate(name0, args0)
+    deps = EmptyDependency(level=self.level, rule_name=None)
+    deps.why = [dep0]
+    for name, args in name_args:
+      dep = Dependency(name, args, None, deps.level)
+      deps.why += [dep.why_me_or_cache(g, None)]
+    return deps
+
+
+def maybe_make_equal_pairs(
+    a: gm.Point,
+    b: gm.Point,
+    c: gm.Point,
+    d: gm.Point,
+    m: gm.Point,
+    n: gm.Point,
+    p: gm.Point,
+    q: gm.Point,
+    ab: gm.Line,
+    mn: gm.Line,
+    g: Any,
+    level: int,
+) -> list[Dependency]:
+  """Make a-b:c-d==m-n:p-q in case a-b==m-n or c-d==p-q."""
+  if ab != mn:
+    return
+  why = []
+  eqname = 'para' if isinstance(ab, gm.Line) else 'cong'
+  colls = [a, b, m, n]
+  if len(set(colls)) > 2 and eqname == 'para':
+    dep = Dependency('collx', colls, None, level)
+    dep.why_me(g, level)
+    why += [dep]
+
+  dep = Dependency(eqname, [c, d, p, q], None, level)
+  dep.why_me(g, level)
+  why += [dep]
+  return why
+
+
+class Dependency(Construction):
+  """Dependency is a predicate that other predicates depend on."""
+
+  def __init__(
+      self, name: str, args: list[gm.Point], rule_name: str, level: int
+  ):
+    super().__init__(name, args)
+    self.rule_name = rule_name or ''
+    self.level = level
+    self.why = []
+
+    self._stat = None
+    self.trace = None
+
+  def _find(self, dep_hashed: tuple[str, ...]) -> Dependency:
+    for w in self.why:
+      f = w._find(dep_hashed)
+      if f:
+        return f
+      if w.hashed() == dep_hashed:
+        return w
+
+  def remove_loop(self) -> Dependency:
+    f = self._find(self.hashed())
+    if f:
+      return f
+    return self
+
+  def copy(self) -> Dependency:
+    dep = Dependency(self.name, self.args, self.rule_name, self.level)
+    dep.trace = self.trace
+    dep.why = list(self.why)
+    return dep
+
+  def why_me_or_cache(self, g: Any, level: int) -> Dependency:
+    if self.hashed() in g.cache:
+      return g.cache[self.hashed()]
+    self.why_me(g, level)
+    return self
+
+  def populate(self, name: str, args: list[gm.Point]) -> Dependency:
+    assert self.rule_name == CONSTRUCTION_RULE, self.rule_name
+    dep = Dependency(self.name, self.args, self.rule_name, self.level)
+    dep.why = list(self.why)
+    return dep
+
+  def why_me(self, g: Any, level: int) -> None:
+    """Figure out the dependencies predicates of self."""
+    name, args = self.name, self.args
+
+    hashed_me = hashed(name, args)
+    if hashed_me in g.cache:
+      dep = g.cache[hashed_me]
+      self.why = dep.why
+      self.rule_name = dep.rule_name
+      return
+
+    if self.name == 'para':
+      a, b, c, d = self.args
+      if {a, b} == {c, d}:
+        self.why = []
+        return
+
+      ab = g._get_line(a, b)
+      cd = g._get_line(c, d)
+      if ab == cd:
+        if {a, b} == {c, d}:
+          self.why = []
+          self.rule_name = ''
+          return
+        dep = Dependency('coll', list({a, b, c, d}), 't??', None)
+        self.why = [dep.why_me_or_cache(g, level)]
+        return
+
+      for (x, y), xy in zip([(a, b), (c, d)], [ab, cd]):
+        x_, y_ = xy.points
+        if {x, y} == {x_, y_}:
+          continue
+        d = Dependency('collx', [x, y, x_, y_], None, level)
+        self.why += [d.why_me_or_cache(g, level)]
+
+      whypara = g.why_equal(ab, cd, None)
+      self.why += whypara
+
+    elif self.name == 'midp':
+      m, a, b = self.args
+      ma = g._get_segment(m, a)
+      mb = g._get_segment(m, b)
+      dep = Dependency('coll', [m, a, b], None, None).why_me_or_cache(g, None)
+      self.why = [dep] + g.why_equal(ma, mb, level)
+
+    elif self.name == 'perp':
+      a, b, c, d = self.args
+      ab = g._get_line(a, b)
+      cd = g._get_line(c, d)
+      for (x, y), xy in zip([(a, b), (c, d)], [ab, cd]):
+        x_, y_ = xy.points
+        if {x, y} == {x_, y_}:
+          continue
+        d = Dependency('collx', [x, y, x_, y_], None, level)
+        self.why += [d.why_me_or_cache(g, level)]
+
+      _, why = why_eqangle(ab._val, cd._val, cd._val, ab._val, level)
+      a, b = ab.points
+      c, d = cd.points
+
+      if hashed(self.name, [a, b, c, d]) != self.hashed():
+        d = Dependency(self.name, [a, b, c, d], None, level)
+        d.why = why
+        why = [d]
+
+      self.why += why
+
+    elif self.name == 'cong':
+      a, b, c, d = self.args
+      ab = g._get_segment(a, b)
+      cd = g._get_segment(c, d)
+
+      self.why = g.why_equal(ab, cd, level)
+
+    elif self.name == 'coll':
+      _, why = gm.line_of_and_why(self.args, level)
+      self.why = why
+
+    elif self.name == 'collx':
+      if g.check_coll(self.args):
+        args = list(set(self.args))
+        hashed_me = hashed('coll', args)
+        if hashed_me in g.cache:
+          dep = g.cache[hashed_me]
+          self.why = [dep]
+          self.rule_name = ''
+          return
+        _, self.why = gm.line_of_and_why(args, level)
+      else:
+        self.name = 'para'
+        self.why_me(g, level)
+
+    elif self.name == 'cyclic':
+      _, why = gm.circle_of_and_why(self.args, level)
+      self.why = why
+
+    elif self.name == 'circle':
+      o, a, b, c = self.args
+      oa = g._get_segment(o, a)
+      ob = g._get_segment(o, b)
+      oc = g._get_segment(o, c)
+      self.why = g.why_equal(oa, ob, level) + g.why_equal(oa, oc, level)
+      
+    elif self.name == 'semicircle':
+      o, a, b, c = self.args  # o: center, a & b: endpoints, c: another point to check
+      oa = g._get_segment(o, a)  # Segment from o to a
+      ob = g._get_segment(o, b)  # Segment from o to b
+      oc = g._get_segment(o, c)  # Segment from o to c
+
+      # Check that segments are equal (radius check)
+      self.why = g.why_equal(oa, ob, level) + g.why_equal(oa, oc, level)
+
+      # Additional checks for semicircle properties can be added here
+      # For example, ensure that point c lies on the semicircle arc defined by a and b
+      self.why += g.why_on_arc(a, b, c, level)  # This function needs to be implemented to check if c is on the arc
+
+
+
+    elif self.name in ['eqangle', 'eqangle6']:
+      a, b, c, d, m, n, p, q = self.args
+
+      ab, why1 = g.get_line_thru_pair_why(a, b)
+      cd, why2 = g.get_line_thru_pair_why(c, d)
+      mn, why3 = g.get_line_thru_pair_why(m, n)
+      pq, why4 = g.get_line_thru_pair_why(p, q)
+
+      if ab is None or cd is None or mn is None or pq is None:
+        if {a, b} == {m, n}:
+          d = Dependency('para', [c, d, p, q], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {a, b} == {c, d}:
+          d = Dependency('para', [p, q, m, n], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {c, d} == {p, q}:
+          d = Dependency('para', [a, b, m, n], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {p, q} == {m, n}:
+          d = Dependency('para', [a, b, c, d], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        return
+
+      for (x, y), xy, whyxy in zip(
+          [(a, b), (c, d), (m, n), (p, q)],
+          [ab, cd, mn, pq],
+          [why1, why2, why3, why4],
+      ):
+        x_, y_ = xy.points
+        if {x, y} == {x_, y_}:
+          continue
+        d = Dependency('collx', [x, y, x_, y_], None, level)
+        d.why = whyxy
+        self.why += [d]
+
+      a, b = ab.points
+      c, d = cd.points
+      m, n = mn.points
+      p, q = pq.points
+      diff = hashed(self.name, [a, b, c, d, m, n, p, q]) != self.hashed()
+
+      whyeqangle = None
+      if ab._val and cd._val and mn._val and pq._val:
+        whyeqangle = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
+
+      if whyeqangle:
+        (dab, dcd, dmn, dpq), whyeqangle = whyeqangle
+        if diff:
+          d = Dependency('eqangle', [a, b, c, d, m, n, p, q], None, level)
+          d.why = whyeqangle
+          whyeqangle = [d]
+        self.why += whyeqangle
+
+      else:
+        if (ab == cd and mn == pq) or (ab == mn and cd == pq):
+          self.why += []
+        elif ab == mn:
+          self.why += maybe_make_equal_pairs(
+              a, b, c, d, m, n, p, q, ab, mn, g, level
+          )
+        elif cd == pq:
+          self.why += maybe_make_equal_pairs(
+              c, d, a, b, p, q, m, n, cd, pq, g, level
+          )
+        elif ab == cd:
+          self.why += maybe_make_equal_pairs(
+              a, b, m, n, c, d, p, q, ab, cd, g, level
+          )
+        elif mn == pq:
+          self.why += maybe_make_equal_pairs(
+              m, n, a, b, p, q, c, d, mn, pq, g, level
+          )
+        elif g.is_equal(ab, mn) or g.is_equal(cd, pq):
+          dep1 = Dependency('para', [a, b, m, n], None, level)
+          dep1.why_me(g, level)
+          dep2 = Dependency('para', [c, d, p, q], None, level)
+          dep2.why_me(g, level)
+          self.why += [dep1, dep2]
+        elif g.is_equal(ab, cd) or g.is_equal(mn, pq):
+          dep1 = Dependency('para', [a, b, c, d], None, level)
+          dep1.why_me(g, level)
+          dep2 = Dependency('para', [m, n, p, q], None, level)
+          dep2.why_me(g, level)
+          self.why += [dep1, dep2]
+        elif ab._val and cd._val and mn._val and pq._val:
+          self.why = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
+
+    elif self.name in ['eqratio', 'eqratio6']:
+      a, b, c, d, m, n, p, q = self.args
+      ab = g._get_segment(a, b)
+      cd = g._get_segment(c, d)
+      mn = g._get_segment(m, n)
+      pq = g._get_segment(p, q)
+
+      if ab is None or cd is None or mn is None or pq is None:
+        if {a, b} == {m, n}:
+          d = Dependency('cong', [c, d, p, q], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {a, b} == {c, d}:
+          d = Dependency('cong', [p, q, m, n], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {c, d} == {p, q}:
+          d = Dependency('cong', [a, b, m, n], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        if {p, q} == {m, n}:
+          d = Dependency('cong', [a, b, c, d], None, level)
+          self.why = [d.why_me_or_cache(g, level)]
+        return
+
+      if ab._val and cd._val and mn._val and pq._val:
+        self.why = why_eqratio(ab._val, cd._val, mn._val, pq._val, level)
+
+      if self.why is None:
+        self.why = []
+        if (ab == cd and mn == pq) or (ab == mn and cd == pq):
+          self.why = []
+        elif ab == mn:
+          self.why += maybe_make_equal_pairs(
+              a, b, c, d, m, n, p, q, ab, mn, g, level
+          )
+        elif cd == pq:
+          self.why += maybe_make_equal_pairs(
+              c, d, a, b, p, q, m, n, cd, pq, g, level
+          )
+        elif ab == cd:
+          self.why += maybe_make_equal_pairs(
+              a, b, m, n, c, d, p, q, ab, cd, g, level
+          )
+        elif mn == pq:
+          self.why += maybe_make_equal_pairs(
+              m, n, a, b, p, q, c, d, mn, pq, g, level
+          )
+        elif g.is_equal(ab, mn) or g.is_equal(cd, pq):
+          dep1 = Dependency('cong', [a, b, m, n], None, level)
+          dep1.why_me(g, level)
+          dep2 = Dependency('cong', [c, d, p, q], None, level)
+          dep2.why_me(g, level)
+          self.why += [dep1, dep2]
+        elif g.is_equal(ab, cd) or g.is_equal(mn, pq):
+          dep1 = Dependency('cong', [a, b, c, d], None, level)
+          dep1.why_me(g, level)
+          dep2 = Dependency('cong', [m, n, p, q], None, level)
+          dep2.why_me(g, level)
+          self.why += [dep1, dep2]
+        elif ab._val and cd._val and mn._val and pq._val:
+          self.why = why_eqangle(ab._val, cd._val, mn._val, pq._val, level)
+
+    elif self.name in ['diff', 'npara', 'nperp', 'ncoll', 'sameside']:
+      self.why = []
+
+    elif self.name == 'simtri':
+      a, b, c, x, y, z = self.args
+      dep1 = Dependency('eqangle', [a, b, a, c, x, y, x, z], '', level)
+      dep1.why_me(g, level)
+      dep2 = Dependency('eqangle', [b, a, b, c, y, x, y, z], '', level)
+      dep2.why_me(g, level)
+      self.rule_name = 'r34'
+      self.why = [dep1, dep2]
+
+    elif self.name == 'contri':
+      a, b, c, x, y, z = self.args
+      dep1 = Dependency('cong', [a, b, x, y], '', level)
+      dep1.why_me(g, level)
+      dep2 = Dependency('cong', [b, c, y, z], '', level)
+      dep2.why_me(g, level)
+      dep3 = Dependency('cong', [c, a, z, x], '', level)
+      dep3.why_me(g, level)
+      self.rule_name = 'r32'
+      self.why = [dep1, dep2, dep3]
+
+    elif self.name == 'ind':
+      pass
+
+    elif self.name == 'aconst':
+      a, b, c, d, ang0 = self.args
+
+      measure = ang0._val
+
+      for ang in measure.neighbors(gm.Angle):
+        if ang == ang0:
+          continue
+        d1, d2 = ang._d
+        l1, l2 = d1._obj, d2._obj
+        (a1, b1), (c1, d1) = l1.points, l2.points
+
+        if not g.check_para_or_coll([a, b, a1, b1]) or not g.check_para_or_coll(
+            [c, d, c1, d1]
+        ):
+          continue
+
+        self.why = []
+        for args in [(a, b, a1, b1), (c, d, c1, d1)]:
+          if g.check_coll(args):
+            if len(set(args)) > 2:
+              dep = Dependency('coll', args, None, None)
+              self.why.append(dep.why_me_or_cache(g, level))
+          else:
+            dep = Dependency('para', args, None, None)
+            self.why.append(dep.why_me_or_cache(g, level))
+
+        self.why += gm.why_equal(ang, ang0)
+        break
+
+    elif self.name == 'rconst':
+      a, b, c, d, rat0 = self.args
+
+      val = rat0._val
+
+      for rat in val.neighbors(gm.Ratio):
+        if rat == rat0:
+          continue
+        l1, l2 = rat._l
+        s1, s2 = l1._obj, l2._obj
+        (a1, b1), (c1, d1) = list(s1.points), list(s2.points)
+
+        if not g.check_cong([a, b, a1, b1]) or not g.check_cong([c, d, c1, d1]):
+          continue
+
+        self.why = []
+        for args in [(a, b, a1, b1), (c, d, c1, d1)]:
+          if len(set(args)) > 2:
+            dep = Dependency('cong', args, None, None)
+            self.why.append(dep.why_me_or_cache(g, level))
+
+        self.why += gm.why_equal(rat, rat0)
+        break
+
+    else:
+      raise ValueError('Not recognize', self.name)
+
+  def hashed(self, rename: bool = False) -> tuple[str, ...]:
+    return hashed(self.name, self.args, rename=rename)
+
+
+def hashed(
+    name: str, args: list[gm.Point], rename: bool = False
+) -> tuple[str, ...]:
+  if name == 's_angle':
+    args = [p.name if not rename else p.new_name for p in args[:-1]] + [
+        str(args[-1])
+    ]
+  else:
+    args = [p.name if not rename else p.new_name for p in args]
+  return hashed_txt(name, args)
+
+
+def hashed_txt(name: str, args: list[str]) -> tuple[str, ...]:
+  """Return a tuple unique to name and args upto arg permutation equivariant."""
+
+  if name in ['const', 'aconst', 'rconst']:
+    a, b, c, d, y = args
+    a, b = sorted([a, b])
+    c, d = sorted([c, d])
+    return name, a, b, c, d, y
+
+  if name in ['npara', 'nperp', 'para', 'cong', 'perp', 'collx']:
+    a, b, c, d = args
+
+    a, b = sorted([a, b])
+    c, d = sorted([c, d])
+    (a, b), (c, d) = sorted([(a, b), (c, d)])
+
+    return (name, a, b, c, d)
+
+  if name in ['midp', 'midpoint']:
+    a, b, c = args
+    b, c = sorted([b, c])
+    return (name, a, b, c)
+
+  if name in ['coll', 'cyclic', 'ncoll', 'diff', 'triangle']:
+    return (name,) + tuple(sorted(list(set(args))))
+
+  if name == 'circle':
+    x, a, b, c = args
+    return (name, x) + tuple(sorted([a, b, c]))
+  
+  if name == 'semicircle':
+    x, a, b, c = args
+    return (name, x) + tuple(sorted([a, b, c]))
+
+  if name in ['eqangle', 'eqratio', 'eqangle6', 'eqratio6']:
+    a, b, c, d, e, f, g, h = args
+    a, b = sorted([a, b])
+    c, d = sorted([c, d])
+    e, f = sorted([e, f])
+    g, h = sorted([g, h])
+    if tuple(sorted([a, b, e, f])) > tuple(sorted([c, d, g, h])):
+      a, b, e, f, c, d, g, h = c, d, g, h, a, b, e, f
+    if (a, b, c, d) > (e, f, g, h):
+      a, b, c, d, e, f, g, h = e, f, g, h, a, b, c, d
+
+    if name == 'eqangle6':
+      name = 'eqangle'
+    if name == 'eqratio6':
+      name = 'eqratio'
+    return (name,) + (a, b, c, d, e, f, g, h)
+
+  if name in ['contri', 'simtri', 'simtri2', 'contri2', 'contri*', 'simtri*']:
+    a, b, c, x, y, z = args
+    (a, x), (b, y), (c, z) = sorted([(a, x), (b, y), (c, z)], key=sorted)
+    (a, b, c), (x, y, z) = sorted([(a, b, c), (x, y, z)], key=sorted)
+    return (name, a, b, c, x, y, z)
+
+  if name in ['eqratio3']:
+    a, b, c, d, o, o = args  # pylint: disable=redeclared-assigned-name
+    (a, c), (b, d) = sorted([(a, c), (b, d)], key=sorted)
+    (a, b), (c, d) = sorted([(a, b), (c, d)], key=sorted)
+    return (name, a, b, c, d, o, o)
+
+  if name in ['sameside', 's_angle']:
+    return (name,) + tuple(args)
+
+  raise ValueError(f'Not recognize {name} to hash.')

ag4masses/alphageometry/rules.txt

@@ -41,3 +41,7 @@ eqratio6 B A B C Q P Q R, eqangle6 B A B C Q P Q R, ncoll A B C => simtri* A B C
 eqratio6 B A B C Q P Q R, eqratio6 C A C B R P R Q, ncoll A B C, cong A B P Q => contri* A B C P Q R
 para a b c d, coll m a d, coll n b c, eqratio6 m a m d n b n c, sameside m a d n b c => para m n a b
 para a b c d, coll m a d, coll n b c, para m n a b => eqratio6 m a m d n b n c
+semicircle O A B C, perp O A A X => eqangle A X A B C A C B
+semicircle O A B C, eqangle A X A B C A C B => perp O A A X
+semicircle O A B C, midp M B C => eqangle A B A C O B O M
+semicircle O A B C, coll M B C, eqangle A B A C O B O M => midp M B C

ag4masses/alphageometry/trace_back.py

@@ -1,374 +1,374 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Implements DAG-level traceback."""
-
-from typing import Any
-
-import geometry as gm
-import pretty as pt
-import problem
-
-
-pretty = pt.pretty
-
-
-def point_levels(
-    setup: list[problem.Dependency], existing_points: list[gm.Point]
-) -> list[tuple[set[gm.Point], list[problem.Dependency]]]:
-  """Reformat setup into levels of point constructions."""
-  levels = []
-  for con in setup:
-    plevel = max([p.plevel for p in con.args if isinstance(p, gm.Point)])
-
-    while len(levels) - 1 < plevel:
-      levels.append((set(), []))
-
-    for p in con.args:
-      if not isinstance(p, gm.Point):
-        continue
-      if existing_points and p in existing_points:
-        continue
-
-      levels[p.plevel][0].add(p)
-
-    cons = levels[plevel][1]
-    cons.append(con)
-
-  return [(p, c) for p, c in levels if p or c]
-
-
-def point_log(
-    setup: list[problem.Dependency],
-    ref_id: dict[tuple[str, ...], int],
-    existing_points=list[gm.Point],
-) -> list[tuple[list[gm.Point], list[problem.Dependency]]]:
-  """Reformat setup into groups of point constructions."""
-  log = []
-
-  levels = point_levels(setup, existing_points)
-
-  for points, cons in levels:
-    for con in cons:
-      if con.hashed() not in ref_id:
-        ref_id[con.hashed()] = len(ref_id)
-
-    log.append((points, cons))
-
-  return log
-
-
-def setup_to_levels(
-    setup: list[problem.Dependency],
-) -> list[list[problem.Dependency]]:
-  """Reformat setup into levels of point constructions."""
-  levels = []
-  for d in setup:
-    plevel = max([p.plevel for p in d.args if isinstance(p, gm.Point)])
-    while len(levels) - 1 < plevel:
-      levels.append([])
-
-    levels[plevel].append(d)
-
-  levels = [lvl for lvl in levels if lvl]
-  return levels
-
-
-def separate_dependency_difference(
-    query: problem.Dependency,
-    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-) -> tuple[
-    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-    list[problem.Dependency],
-    list[problem.Dependency],
-    set[gm.Point],
-    set[gm.Point],
-]:
-  """Identify and separate the dependency difference."""
-  setup = []
-  log_, log = log, []
-  for prems, cons in log_:
-    if not prems:
-      setup.extend(cons)
-      continue
-    cons_ = []
-    for con in cons:
-      if con.rule_name == 'c0':
-        setup.append(con)
-      else:
-        cons_.append(con)
-    if not cons_:
-      continue
-
-    prems = [p for p in prems if p.name != 'ind']
-    log.append((prems, cons_))
-
-  points = set(query.args)
-  queue = list(query.args)
-  i = 0
-  while i < len(queue):
-    q = queue[i]
-    i += 1
-    if not isinstance(q, gm.Point):
-      continue
-    for p in q.rely_on:
-      if p not in points:
-        points.add(p)
-        queue.append(p)
-
-  setup_, setup, aux_setup, aux_points = setup, [], [], set()
-  for con in setup_:
-    if con.name == 'ind':
-      continue
-    elif any([p not in points for p in con.args if isinstance(p, gm.Point)]):
-      aux_setup.append(con)
-      aux_points.update(
-          [p for p in con.args if isinstance(p, gm.Point) and p not in points]
-      )
-    else:
-      setup.append(con)
-
-  return log, setup, aux_setup, points, aux_points
-
-
-def recursive_traceback(
-    query: problem.Dependency,
-) -> list[tuple[list[problem.Dependency], list[problem.Dependency]]]:
-  """Recursively traceback from the query, i.e. the conclusion."""
-  visited = set()
-  log = []
-  stack = []
-
-  def read(q: problem.Dependency) -> None:
-    q = q.remove_loop()
-    hashed = q.hashed()
-    if hashed in visited:
-      return
-
-    if hashed[0] in ['ncoll', 'npara', 'nperp', 'diff', 'sameside']:
-      return
-
-    nonlocal stack
-
-    stack.append(hashed)
-    prems = []
-
-    if q.rule_name != problem.CONSTRUCTION_RULE:
-      all_deps = []
-      dep_names = set()
-      for d in q.why:
-        if d.hashed() in dep_names:
-          continue
-        dep_names.add(d.hashed())
-        all_deps.append(d)
-
-      for d in all_deps:
-        h = d.hashed()
-        if h not in visited:
-          read(d)
-        if h in visited:
-          prems.append(d)
-
-    visited.add(hashed)
-    hashs = sorted([d.hashed() for d in prems])
-    found = False
-    for ps, qs in log:
-      if sorted([d.hashed() for d in ps]) == hashs:
-        qs += [q]
-        found = True
-        break
-    if not found:
-      log.append((prems, [q]))
-
-    stack.pop(-1)
-
-  read(query)
-
-  # post process log: separate multi-conclusion lines
-  log_, log = log, []
-  for ps, qs in log_:
-    for q in qs:
-      log.append((ps, [q]))
-
-  return log
-
-
-def collx_to_coll_setup(
-    setup: list[problem.Dependency],
-) -> list[problem.Dependency]:
-  """Convert collx to coll in setups."""
-  result = []
-  for level in setup_to_levels(setup):
-    hashs = set()
-    for dep in level:
-      if dep.name == 'collx':
-        dep.name = 'coll'
-        dep.args = list(set(dep.args))
-
-      if dep.hashed() in hashs:
-        continue
-      hashs.add(dep.hashed())
-      result.append(dep)
-
-  return result
-
-
-def collx_to_coll(
-    setup: list[problem.Dependency],
-    aux_setup: list[problem.Dependency],
-    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-) -> tuple[
-    list[problem.Dependency],
-    list[problem.Dependency],
-    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-]:
-  """Convert collx to coll and dedup."""
-  setup = collx_to_coll_setup(setup)
-  aux_setup = collx_to_coll_setup(aux_setup)
-
-  con_set = set([p.hashed() for p in setup + aux_setup])
-  log_, log = log, []
-  for prems, cons in log_:
-    prem_set = set()
-    prems_, prems = prems, []
-    for p in prems_:
-      if p.name == 'collx':
-        p.name = 'coll'
-        p.args = list(set(p.args))
-      if p.hashed() in prem_set:
-        continue
-      prem_set.add(p.hashed())
-      prems.append(p)
-
-    cons_, cons = cons, []
-    for c in cons_:
-      if c.name == 'collx':
-        c.name = 'coll'
-        c.args = list(set(c.args))
-      if c.hashed() in con_set:
-        continue
-      con_set.add(c.hashed())
-      cons.append(c)
-
-    if not cons or not prems:
-      continue
-
-    log.append((prems, cons))
-
-  return setup, aux_setup, log
-
-
-def get_logs(
-    query: problem.Dependency, g: Any, merge_trivials: bool = False
-) -> tuple[
-    list[problem.Dependency],
-    list[problem.Dependency],
-    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-    set[gm.Point],
-]:
-  """Given a DAG and conclusion N, return the premise, aux, proof."""
-  query = query.why_me_or_cache(g, query.level)
-  log = recursive_traceback(query)
-  log, setup, aux_setup, setup_points, _ = separate_dependency_difference(
-      query, log
-  )
-
-  setup, aux_setup, log = collx_to_coll(setup, aux_setup, log)
-
-  setup, aux_setup, log = shorten_and_shave(
-      setup, aux_setup, log, merge_trivials
-  )
-
-  return setup, aux_setup, log, setup_points
-
-
-def shorten_and_shave(
-    setup: list[problem.Dependency],
-    aux_setup: list[problem.Dependency],
-    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-    merge_trivials: bool = False,
-) -> tuple[
-    list[problem.Dependency],
-    list[problem.Dependency],
-    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-]:
-  """Shorten the proof by removing unused predicates."""
-  log, _ = shorten_proof(log, merge_trivials=merge_trivials)
-
-  all_prems = sum([list(prems) for prems, _ in log], [])
-  all_prems = set([p.hashed() for p in all_prems])
-  setup = [d for d in setup if d.hashed() in all_prems]
-  aux_setup = [d for d in aux_setup if d.hashed() in all_prems]
-  return setup, aux_setup, log
-
-
-def join_prems(
-    con: problem.Dependency,
-    con2prems: dict[tuple[str, ...], list[problem.Dependency]],
-    expanded: set[tuple[str, ...]],
-) -> list[problem.Dependency]:
-  """Join proof steps with the same premises."""
-  h = con.hashed()
-  if h in expanded or h not in con2prems:
-    return [con]
-
-  result = []
-  for p in con2prems[h]:
-    result += join_prems(p, con2prems, expanded)
-  return result
-
-
-def shorten_proof(
-    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-    merge_trivials: bool = False,
-) -> tuple[
-    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
-    dict[tuple[str, ...], list[problem.Dependency]],
-]:
-  """Join multiple trivials proof steps into one."""
-  pops = set()
-  con2prem = {}
-  for prems, cons in log:
-    assert len(cons) == 1
-    con = cons[0]
-    if con.rule_name == '':  # pylint: disable=g-explicit-bool-comparison
-      con2prem[con.hashed()] = prems
-    elif not merge_trivials:
-      # except for the ones that are premises to non-trivial steps.
-      pops.update({p.hashed() for p in prems})
-
-  for p in pops:
-    if p in con2prem:
-      con2prem.pop(p)
-
-  expanded = set()
-  log2 = []
-  for i, (prems, cons) in enumerate(log):
-    con = cons[0]
-    if i < len(log) - 1 and con.hashed() in con2prem:
-      continue
-
-    hashs = set()
-    new_prems = []
-
-    for p in sum([join_prems(p, con2prem, expanded) for p in prems], []):
-      if p.hashed() not in hashs:
-        new_prems.append(p)
-        hashs.add(p.hashed())
-
-    log2 += [(new_prems, [con])]
-    expanded.add(con.hashed())
-
-  return log2, con2prem
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Implements DAG-level traceback."""
+
+from typing import Any
+
+import geometry as gm
+import pretty as pt
+import problem
+
+
+pretty = pt.pretty
+
+
+def point_levels(
+    setup: list[problem.Dependency], existing_points: list[gm.Point]
+) -> list[tuple[set[gm.Point], list[problem.Dependency]]]:
+  """Reformat setup into levels of point constructions."""
+  levels = []
+  for con in setup:
+    plevel = max([p.plevel for p in con.args if isinstance(p, gm.Point)])
+
+    while len(levels) - 1 < plevel:
+      levels.append((set(), []))
+
+    for p in con.args:
+      if not isinstance(p, gm.Point):
+        continue
+      if existing_points and p in existing_points:
+        continue
+
+      levels[p.plevel][0].add(p)
+
+    cons = levels[plevel][1]
+    cons.append(con)
+
+  return [(p, c) for p, c in levels if p or c]
+
+
+def point_log(
+    setup: list[problem.Dependency],
+    ref_id: dict[tuple[str, ...], int],
+    existing_points=list[gm.Point],
+) -> list[tuple[list[gm.Point], list[problem.Dependency]]]:
+  """Reformat setup into groups of point constructions."""
+  log = []
+
+  levels = point_levels(setup, existing_points)
+
+  for points, cons in levels:
+    for con in cons:
+      if con.hashed() not in ref_id:
+        ref_id[con.hashed()] = len(ref_id)
+
+    log.append((points, cons))
+
+  return log
+
+
+def setup_to_levels(
+    setup: list[problem.Dependency],
+) -> list[list[problem.Dependency]]:
+  """Reformat setup into levels of point constructions."""
+  levels = []
+  for d in setup:
+    plevel = max([p.plevel for p in d.args if isinstance(p, gm.Point)])
+    while len(levels) - 1 < plevel:
+      levels.append([])
+
+    levels[plevel].append(d)
+
+  levels = [lvl for lvl in levels if lvl]
+  return levels
+
+
+def separate_dependency_difference(
+    query: problem.Dependency,
+    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+) -> tuple[
+    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+    list[problem.Dependency],
+    list[problem.Dependency],
+    set[gm.Point],
+    set[gm.Point],
+]:
+  """Identify and separate the dependency difference."""
+  setup = []
+  log_, log = log, []
+  for prems, cons in log_:
+    if not prems:
+      setup.extend(cons)
+      continue
+    cons_ = []
+    for con in cons:
+      if con.rule_name == 'c0':
+        setup.append(con)
+      else:
+        cons_.append(con)
+    if not cons_:
+      continue
+
+    prems = [p for p in prems if p.name != 'ind']
+    log.append((prems, cons_))
+
+  points = set(query.args)
+  queue = list(query.args)
+  i = 0
+  while i < len(queue):
+    q = queue[i]
+    i += 1
+    if not isinstance(q, gm.Point):
+      continue
+    for p in q.rely_on:
+      if p not in points:
+        points.add(p)
+        queue.append(p)
+
+  setup_, setup, aux_setup, aux_points = setup, [], [], set()
+  for con in setup_:
+    if con.name == 'ind':
+      continue
+    elif any([p not in points for p in con.args if isinstance(p, gm.Point)]):
+      aux_setup.append(con)
+      aux_points.update(
+          [p for p in con.args if isinstance(p, gm.Point) and p not in points]
+      )
+    else:
+      setup.append(con)
+
+  return log, setup, aux_setup, points, aux_points
+
+
+def recursive_traceback(
+    query: problem.Dependency,
+) -> list[tuple[list[problem.Dependency], list[problem.Dependency]]]:
+  """Recursively traceback from the query, i.e. the conclusion."""
+  visited = set()
+  log = []
+  stack = []
+
+  def read(q: problem.Dependency) -> None:
+    q = q.remove_loop()
+    hashed = q.hashed()
+    if hashed in visited:
+      return
+
+    if hashed[0] in ['ncoll', 'npara', 'nperp', 'diff', 'sameside']:
+      return
+
+    nonlocal stack
+
+    stack.append(hashed)
+    prems = []
+
+    if q.rule_name != problem.CONSTRUCTION_RULE:
+      all_deps = []
+      dep_names = set()
+      for d in q.why:
+        if d.hashed() in dep_names:
+          continue
+        dep_names.add(d.hashed())
+        all_deps.append(d)
+
+      for d in all_deps:
+        h = d.hashed()
+        if h not in visited:
+          read(d)
+        if h in visited:
+          prems.append(d)
+
+    visited.add(hashed)
+    hashs = sorted([d.hashed() for d in prems])
+    found = False
+    for ps, qs in log:
+      if sorted([d.hashed() for d in ps]) == hashs:
+        qs += [q]
+        found = True
+        break
+    if not found:
+      log.append((prems, [q]))
+
+    stack.pop(-1)
+
+  read(query)
+
+  # post process log: separate multi-conclusion lines
+  log_, log = log, []
+  for ps, qs in log_:
+    for q in qs:
+      log.append((ps, [q]))
+
+  return log
+
+
+def collx_to_coll_setup(
+    setup: list[problem.Dependency],
+) -> list[problem.Dependency]:
+  """Convert collx to coll in setups."""
+  result = []
+  for level in setup_to_levels(setup):
+    hashs = set()
+    for dep in level:
+      if dep.name == 'collx':
+        dep.name = 'coll'
+        dep.args = list(set(dep.args))
+
+      if dep.hashed() in hashs:
+        continue
+      hashs.add(dep.hashed())
+      result.append(dep)
+
+  return result
+
+
+def collx_to_coll(
+    setup: list[problem.Dependency],
+    aux_setup: list[problem.Dependency],
+    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+) -> tuple[
+    list[problem.Dependency],
+    list[problem.Dependency],
+    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+]:
+  """Convert collx to coll and dedup."""
+  setup = collx_to_coll_setup(setup)
+  aux_setup = collx_to_coll_setup(aux_setup)
+
+  con_set = set([p.hashed() for p in setup + aux_setup])
+  log_, log = log, []
+  for prems, cons in log_:
+    prem_set = set()
+    prems_, prems = prems, []
+    for p in prems_:
+      if p.name == 'collx':
+        p.name = 'coll'
+        p.args = list(set(p.args))
+      if p.hashed() in prem_set:
+        continue
+      prem_set.add(p.hashed())
+      prems.append(p)
+
+    cons_, cons = cons, []
+    for c in cons_:
+      if c.name == 'collx':
+        c.name = 'coll'
+        c.args = list(set(c.args))
+      if c.hashed() in con_set:
+        continue
+      con_set.add(c.hashed())
+      cons.append(c)
+
+    if not cons or not prems:
+      continue
+
+    log.append((prems, cons))
+
+  return setup, aux_setup, log
+
+
+def get_logs(
+    query: problem.Dependency, g: Any, merge_trivials: bool = False
+) -> tuple[
+    list[problem.Dependency],
+    list[problem.Dependency],
+    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+    set[gm.Point],
+]:
+  """Given a DAG and conclusion N, return the premise, aux, proof."""
+  query = query.why_me_or_cache(g, query.level)
+  log = recursive_traceback(query)
+  log, setup, aux_setup, setup_points, _ = separate_dependency_difference(
+      query, log
+  )
+
+  setup, aux_setup, log = collx_to_coll(setup, aux_setup, log)
+
+  setup, aux_setup, log = shorten_and_shave(
+      setup, aux_setup, log, merge_trivials
+  )
+
+  return setup, aux_setup, log, setup_points
+
+
+def shorten_and_shave(
+    setup: list[problem.Dependency],
+    aux_setup: list[problem.Dependency],
+    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+    merge_trivials: bool = False,
+) -> tuple[
+    list[problem.Dependency],
+    list[problem.Dependency],
+    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+]:
+  """Shorten the proof by removing unused predicates."""
+  log, _ = shorten_proof(log, merge_trivials=merge_trivials)
+
+  all_prems = sum([list(prems) for prems, _ in log], [])
+  all_prems = set([p.hashed() for p in all_prems])
+  setup = [d for d in setup if d.hashed() in all_prems]
+  aux_setup = [d for d in aux_setup if d.hashed() in all_prems]
+  return setup, aux_setup, log
+
+
+def join_prems(
+    con: problem.Dependency,
+    con2prems: dict[tuple[str, ...], list[problem.Dependency]],
+    expanded: set[tuple[str, ...]],
+) -> list[problem.Dependency]:
+  """Join proof steps with the same premises."""
+  h = con.hashed()
+  if h in expanded or h not in con2prems:
+    return [con]
+
+  result = []
+  for p in con2prems[h]:
+    result += join_prems(p, con2prems, expanded)
+  return result
+
+
+def shorten_proof(
+    log: list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+    merge_trivials: bool = False,
+) -> tuple[
+    list[tuple[list[problem.Dependency], list[problem.Dependency]]],
+    dict[tuple[str, ...], list[problem.Dependency]],
+]:
+  """Join multiple trivials proof steps into one."""
+  pops = set()
+  con2prem = {}
+  for prems, cons in log:
+    assert len(cons) == 1
+    con = cons[0]
+    if con.rule_name == '':  # pylint: disable=g-explicit-bool-comparison
+      con2prem[con.hashed()] = prems
+    elif not merge_trivials:
+      # except for the ones that are premises to non-trivial steps.
+      pops.update({p.hashed() for p in prems})
+
+  for p in pops:
+    if p in con2prem:
+      con2prem.pop(p)
+
+  expanded = set()
+  log2 = []
+  for i, (prems, cons) in enumerate(log):
+    con = cons[0]
+    if i < len(log) - 1 and con.hashed() in con2prem:
+      continue
+
+    hashs = set()
+    new_prems = []
+
+    for p in sum([join_prems(p, con2prem, expanded) for p in prems], []):
+      if p.hashed() not in hashs:
+        new_prems.append(p)
+        hashs.add(p.hashed())
+
+    log2 += [(new_prems, [con])]
+    expanded.add(con.hashed())
+
+  return log2, con2prem

ag4masses/alphageometry/transformer_layer.py

@@ -1,527 +1,526 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""A single transformer layer in inference mode.
-
-Modified
-https://github.com/google-research/meliad/blob/main/transformer/transformer_layer.py
-To accommodate sequence packing + kv cache + relative position during test time.
-"""
-
-from typing import Callable, Mapping, NewType, Optional, Tuple
-
-from absl import logging
-import gin
-import jax
-import jax.numpy as jnp
-from transformer import attention
-from transformer import nn_components
-from transformer import position
-from transformer import transformer_layer
-
-
-Array = jnp.ndarray
-DecoderState = NewType("DecoderState", Mapping[str, Array])
-WindowState = Optional[Tuple[attention.KVITuple, Array]]
-
-
-@jax.vmap
-def update_slice_in_dim_1(array: Array, update: Array, idx: Array) -> Array:
-  """Update a stored keys/values slice for different-lengthed seqs in batch."""
-  return jax.lax.dynamic_update_slice_in_dim(array, update, idx, axis=0)
-
-
-def slice_in_dim_1(window_length: int) -> Callable[[Array, Array], Array]:
-  @jax.vmap
-  def fn(array: Array, idx: Array) -> Array:
-    return jax.lax.dynamic_slice_in_dim(array, idx, window_length, axis=0)
-
-  return fn
-
-
-@gin.configurable
-class TransformerLayerGenerate(transformer_layer.TransformerLayer):
-  """Full transformer layer, with attention."""
-
-  def _next_decoder_state(
-      self, decoder_state: DecoderState, keys: Array, values: Array
-  ) -> Tuple[DecoderState, Array, Array]:
-    """Compute the next decoder state, and return keys,values to attend to.
-
-    The keys,values returned from this function are drawn from the prior
-    decoding state, and comprise a full window of local context.
-
-    Args:
-      decoder_state: The current decoder state, initially created using
-        init_decoder_state().
-      keys: The key for the current token, of shape (batch_size, 1, dim)
-      values: The value for the current token of shape (batch_size, 1, dim)
-
-    Returns:
-      (next_decoder_state,
-       window of keys of shape (batch_size, window_length, dim),
-       window of values of shape (batch_size, window_length, dim))
-    """
-
-    assert keys.shape[1] == 1  # single-token autoregressive decoding.
-
-    # Unpack decoder_state
-    stored_keys = decoder_state["keys"]
-    stored_values = decoder_state["values"]
-    curr_index = decoder_state["current_index"]
-
-    # Slice to get window_length-sized chunk of previous keys,values.
-    out_decoder_state = {}
-    curr_win_index = curr_index - self.window_length
-
-    # out_keys = jax.lax.dynamic_slice_in_dim(
-    #     stored_keys, curr_win_index, self.window_length, axis=1)
-    out_keys = slice_in_dim_1(self.window_length)(stored_keys, curr_win_index)
-
-    # out_values = jax.lax.dynamic_slice_in_dim(
-    #     stored_values, curr_win_index, self.window_length, axis=1)
-    out_values = slice_in_dim_1(self.window_length)(
-        stored_values, curr_win_index
-    )
-
-    # Write current keys,values to stored keys, values.
-    # stored_keys = jax.lax.dynamic_update_slice_in_dim(
-    #     stored_keys, keys, curr_index, axis=1)
-    stored_keys = update_slice_in_dim_1(stored_keys, keys, curr_index)
-    # stored_values = jax.lax.dynamic_update_slice_in_dim(
-    #     stored_values, values, curr_index, axis=1)
-    stored_values = update_slice_in_dim_1(stored_values, values, curr_index)
-    curr_index = curr_index + 1
-
-    # Pack a new decoder_state object.
-    out_decoder_state["keys"] = stored_keys
-    out_decoder_state["values"] = stored_values
-    out_decoder_state["current_index"] = curr_index
-    out_decoder_state["relative_position_bias"] = decoder_state[
-        "relative_position_bias"
-    ]
-    out_decoder_state["recurrent_kvq"] = decoder_state["recurrent_kvq"]
-
-    return (DecoderState(out_decoder_state), out_keys, out_values)
-
-  def __call__(
-      self,
-      xs: Array,
-      start_of_sequence: Array,
-      *,
-      importance: Optional[Array] = None,
-      cross_attention_kv: Optional[Tuple[Array, Array]] = None,
-      window_state: Optional[WindowState] = None,
-      decoder_state: Optional[DecoderState] = None,
-  ):
-    """Computes attention over a sequence of inputs.
-
-    Args:
-      xs: input sequence of shape (batch_size, sequence_length, num_hidden)
-      start_of_sequence: An input array of shape (batch_size)  --- The following
-        must be passed by keyword only. ---
-      importance: Array of shape (batch_size, sequence_length). An importance
-        bias for attention.
-      cross_attention_kv: Keys and values from encoder for cross-attention.
-      window_state: State object which contains context from the prior window
-        when using a transformer-XL or sliding window. Initially created with
-        load_window_state().
-      decoder_state: State object for autoregressive decoding, initially created
-        with from init_decoder_state().
-
-    Returns:
-      (ys: outputs of shape (batch_size, sequence_length, num_hidden),
-       importance_score: importance score for the next layer,
-       next_window_state: state to pass to the next window,
-       next_decoder_state: next decoder state for autoregressive decoding,
-       viz_dict: dictionary of visualizations
-      )
-    """
-
-    xs = jnp.asarray(xs, dtype=self.dtype)
-    logging.info("tlayer: recurrent = %r", self.recurrent_attention)
-    logging.info("tlayer: compute_importance = %r", self.compute_importance)
-
-    is_training = self.mode == "train"
-
-    # Compute keys, values and queries.
-    # ---------------------------------
-    logging.info("tlayer: compute keys,values,queries.")
-    (keys, values, queries, queries2) = self.tbase.kvq(xs)
-    attention_scale_factors = self.tbase.attention_scale_factors()
-    (_, sequence_length, num_heads, _) = queries.shape  # (b, k, h, d)
-
-    # Get biases and masks that are shared across windows.
-    # ----------------------------------------------------
-    if decoder_state is not None:
-      logging.info("tlayer: using autoregressive decoder.")
-      # When decoding, prior keys,values are loaded from the decoder state.
-      # Other values are precomputed, and loaded from the decoder state.
-      # The decoder state will be updated with the current token.
-      assert window_state is None
-
-      prev_kvi = None
-      recurrent_state = None  # Use precomputed recurrent_kvq.
-      cross_attention_kv = None
-      rel_position_bias = decoder_state["relative_position_bias"]
-      causal_mask = None
-      dropout_multiplier = None
-
-      # Reuse cached recurrent keys,values for each token.
-      cached_recurrent_kvq = decoder_state["recurrent_kvq"]
-      if cached_recurrent_kvq is not None:
-        assert cross_attention_kv is None
-        cross_attention_kv = (cached_recurrent_kvq[0], cached_recurrent_kvq[1])
-      del cached_recurrent_kvq
-
-      # Get a full window of keys,values and update decoder state.
-      (decoder_state, keys, values) = self._next_decoder_state(
-          decoder_state, keys, values
-      )
-
-      # Each query attends to window_length prior keys.
-      assert keys.shape[1] == self.window_length
-      kq_relative_offset = self.window_length
-
-      if not self.use_long_xl_architecture:
-        kqpos = position.relative_positions(
-            1, self.window_length, offset=0
-        )  # 2D mask
-        current_idx = decoder_state["current_index"]
-
-        # add (batch, heads) dims for kqpos
-        kqpos = jnp.expand_dims(kqpos, axis=(0, 1))
-        kqpos = jnp.tile(kqpos, (1, self.num_heads, 1, 1))
-
-        # add (_, heads, _) dim for current_idx
-        current_idx = jnp.expand_dims(current_idx, axis=(1, 2, 3))
-
-        causal_mask = kqpos > self.window_length * 2 - current_idx
-    else:
-      logging.info("tlayer: windowed attention.")
-      # When training, attention is done using windows or chunks, and prior
-      # context (e.g. keys,values from the previous window) is stored in the
-      # window_state object.
-      (prev_kvi, recurrent_state) = (
-          window_state  # pytype: disable=attribute-error
-      )
-
-      # Get the size of the sliding window for pos bias, dropout, & causal mask.
-      (num_queries, num_keys) = attention.sliding_attention_window_shape(
-          (keys, values, importance),
-          prev_kvi,
-          queries,
-          window_length=self.window_length,
-      )
-      kq_relative_offset = num_keys - num_queries
-
-      # Get the relative position bias.
-      # The bias doesn't depend on the query content, and so can be precomputed.
-      if self.relative_positions is not None:
-        rel_position_bias = self.relative_positions(
-            num_queries, num_keys, bidirectional=False
-        )
-      else:
-        rel_position_bias = None
-
-      # Get causal mask.
-      if self.use_causal_mask:
-        causal_mask = position.causal_mask(
-            num_queries, num_keys, window_length=self.window_length
-        )
-      else:
-        causal_mask = None
-
-      # Apply dropout to the attention matrix.
-      # The mask will be broadcast across batches and windows.
-      if self.attn_dropout_rate > 0.0 and is_training:
-        dropout_rng = self.make_rng("dropout")
-        attn_shape = (self.num_heads, num_queries, num_keys)
-        dropout_multiplier = nn_components.dropout_multiplier_mask(
-            dropout_rng, self.attn_dropout_rate, attn_shape, self.dtype
-        )
-      else:
-        dropout_multiplier = None
-
-    # Load and store values into external memory, if memory is not None.
-    # ------------------------------------------------------------------
-    (mode, _, update_memory) = self._get_cache_name_from_mode(self.mode)
-    external_kv = self._query_external_memory(
-        keys,
-        values,
-        queries,
-        start_of_sequence=start_of_sequence,
-        mode=mode,
-        update_memory=decoder_state is None and update_memory,
-    )
-
-    if (
-        self.memory is not None
-        and self.memory_combine_with_local == "TRAINABLE_WEIGHTED_MEAN"
-    ):
-      external_memory_bias = jnp.asarray(self.memory_bias, dtype=self.dtype)
-      external_memory_bias = jnp.reshape(
-          external_memory_bias, (1, 1, num_heads, 1)
-      )
-      external_memory_bias = jax.nn.sigmoid(external_memory_bias)
-    else:
-      external_memory_bias = None
-
-    # Compute the number of windows.
-    # ------------------------------
-    if sequence_length < self.window_length:
-      num_windows = 1  # Happens with autoregressive decoding.
-    elif sequence_length == self.window_length:
-      num_windows = 1
-      if self.use_long_xl_architecture:
-        assert prev_kvi is not None
-    else:
-      if not self.use_long_xl_architecture:
-        raise ValueError("Can only use sliding window with Transformer XL.")
-      num_windows = sequence_length // self.window_length
-      if (num_windows * self.window_length) != sequence_length:
-        raise ValueError(
-            f"Window length {self.window_length} must be a "
-            + f"multiple of sequence length {sequence_length}"
-        )
-    logging.info("tlayer: num_windows = %d.", num_windows)
-
-    # Define the function to do attention within a single window.
-    # ---------------------------------------------------------
-    def single_window_attention(
-        carry: tuple[Array, Array], inputs_w: tuple[Array, Array]
-    ) -> tuple[tuple[Array, Array], tuple[Array, Array]]:
-      # This function uses the following variables from the outer scope.
-      # They are listed here for clarity.
-      nonlocal rel_position_bias
-      nonlocal causal_mask
-      nonlocal kq_relative_offset
-      nonlocal dropout_multiplier
-      nonlocal attention_scale_factors
-      nonlocal external_memory_bias
-      nonlocal cross_attention_kv  # externally supplied.
-
-      # keys,values,queries over the whole sequence will be split into chunks.
-      # xs_w, kvqi_w, etc. are the chunk for the current window.
-      (prev_kvi_w, rec_state) = carry  # carried from one window to the next.
-      (kvqi_w, external_kv_w) = inputs_w  # inputs to the current window.
-      # (keys_curr_w, values_curr_w, _, _, importance_curr_w) = kvqi_w
-
-      # Concatenate keys,values from the previous window with the current
-      # window to implement sliding window attention.
-      (kvqi_w, next_kvi_w) = attention.concat_kvqi(kvqi_w, prev_kvi_w)
-      (keys_w, values_w, queries_w, queries2_w, importance_w) = kvqi_w
-
-      # Perform recurrent attention within the current window to get the next
-      # recurrent state, and set up cross attention.
-      if rec_state is not None:
-        logging.info("tlayer: recurrent attention.")
-
-        # NOTE -- recurrent states and input tokens are handled separately,
-        # because they have separate learned positional embeddings.  Due to
-        # the way TransformerBase does cross-attention, this means that we use
-        # separate key,value layers for rec_state and tokens_w.
-
-        # Keys, values, queries from recurrent state.
-        logging.info("tlayer: recurrent kvq.")
-        rec_kvq = self.recurrent_tbase.kvq(rec_state)
-        r_scale_factors = self.recurrent_tbase.attention_scale_factors()
-        (r_keys, r_values, r_queries, r_queries2) = rec_kvq
-
-        # Joint attention over both recurrent states and input tokens.
-        logging.info("tlayer: recurrent self-attention.")
-        r_attn_ys = attention.simple_attention(
-            r_keys,
-            r_values,
-            r_queries,
-            None,
-            scale_factor=r_scale_factors[0],
-            dtype=self.dtype,
-        )
-
-        logging.info("tlayer: recurrent cross-attention.")
-        r_cross_attn_ys = attention.simple_attention(
-            keys_w,
-            values_w,
-            r_queries2,
-            importance_w,
-            scale_factor=r_scale_factors[1],
-            dtype=self.dtype,
-        )
-
-        # Recurrent post-attention FFN.
-        logging.info("tlayer: recurrent ffn.")
-        next_rec_state = self.recurrent_tbase.post_attn_ffn(
-            rec_state, r_attn_ys, r_cross_attn_ys
-        )
-
-        # Get keys and values for cross-attention from recurrent state.
-        assert cross_attention_kv is None
-        local_cross_attention_kv = (r_keys, r_values)
-      else:
-        # Get keys and values for cross-attention from external argument.
-        next_rec_state = None
-        local_cross_attention_kv = cross_attention_kv
-
-      # If using RoPE, keys and queries are rotated before self-attention.
-      if self.relative_position_type == "rotary":
-        logging.info(
-            "Using rotary position encodings (RoPE), offset = %d",
-            kq_relative_offset,
-        )
-        (keys_w, queries_w) = position.rotate_kq(
-            keys_w, queries_w, max_wavelength=10_000, offset=kq_relative_offset
-        )
-
-      # Self-attention over input tokens.
-      logging.info("tlayer: self-attention.")
-      attn_ys_w = attention.simple_attention(
-          keys_w,
-          values_w,
-          queries_w,
-          importance_w,
-          relative_position_bias=rel_position_bias,
-          scale_factor=attention_scale_factors[0],
-          causal_mask=causal_mask,
-          dropout_multiplier=dropout_multiplier,
-          dtype=self.dtype,
-      )
-
-      # Attention over external memory.
-      if external_kv_w is not None:
-        (external_keys_w, external_values_w) = external_kv_w
-        y_ext = attention.external_attention(
-            external_keys_w,
-            external_values_w,
-            queries_w,
-            scale_factor=attention_scale_factors[0],
-        )
-        if external_memory_bias is not None:
-          ebias = external_memory_bias
-          attn_ys_w = (attn_ys_w * (1 - ebias)) + (y_ext * ebias)
-        elif self.memory_combine_with_local == "ADD":
-          attn_ys_w += y_ext
-        elif self.memory_combine_with_local == "STOP_FORWARD":
-          attn_ys_w = y_ext + (attn_ys_w - jax.lax.stop_gradient(attn_ys_w))
-        else:
-          raise ValueError(
-              f"Unexpected setting: {self.memory_combine_with_local = }"
-          )
-
-      # Cross attention from input tokens to encoder or recurrent state.
-      if local_cross_attention_kv is not None:
-        logging.info("tlayer: cross-attention.")
-        (c_keys, c_values) = local_cross_attention_kv
-
-        # Cross-attention using queries2.
-        cross_attn_ys_w = attention.simple_attention(
-            c_keys,
-            c_values,
-            queries2_w,
-            None,
-            scale_factor=attention_scale_factors[1],
-            dtype=self.dtype,
-        )
-      else:
-        cross_attn_ys_w = None
-
-      # End function single_window_attention(...)
-      return ((next_kvi_w, next_rec_state), (attn_ys_w, cross_attn_ys_w))
-
-    # Initialize recurrent_tbase before calling jax.lax.scan.
-    # Otherwise flax will throw a tantrum.
-    if (
-        self.recurrent_attention
-        and 0 <= self.max_unrolled_windows
-        and self.max_unrolled_windows < num_windows
-    ):
-      logging.info("tlayer: force initialization of recurrent_tbase.")
-      self.recurrent_tbase.force_init(recurrent_state)
-
-    # Perform sliding window attention over all keys,values,queries.
-    # --------------------------------------------------------------
-    initial_carry = (prev_kvi, recurrent_state)  # window state.
-    kvqi = (keys, values, queries, queries2, importance)
-    attn_inputs = (kvqi, external_kv)
-    (next_carry, attn_outputs) = attention.split_and_scan(
-        single_window_attention,
-        initial_carry,
-        attn_inputs,
-        sections=num_windows,
-        axis=1,
-        max_unrolled_windows=self.max_unrolled_windows,
-    )
-    (attn_ys, cross_attn_ys) = attn_outputs
-
-    logging.info("tlayer: End windows.")
-
-    # Post-attention MLP, resnet, and FFN.
-    # ------------------------------------
-    logging.info("tlayer: final FFN.")
-    ys = self.tbase.post_attn_ffn(xs, attn_ys, cross_attn_ys)
-
-    # Compute importance scores for each token if requested.
-    if self.compute_importance:
-      (batch_size, sequence_length, _) = ys.shape
-      importance_score = self.importance_layer(ys)
-      importance_score = importance_score.reshape((batch_size, sequence_length))
-    else:
-      importance_score = None
-
-    next_window_state = next_carry if window_state is not None else None
-    viz_dict = {}  # Visualizations, not currently enabled.
-    return (ys, importance_score, next_window_state, decoder_state, viz_dict)
-
-  def init_decoder_state_vanilla(
-      self, sequence_length: int, start_of_sequence: Array
-  ) -> DecoderState:
-    """Initialize decoder state for autoregressive generation.
-
-    Args:
-      sequence_length: The maximum length of the sequence to generate.
-      start_of_sequence: Array of boolean of shape (batch_size,) True if
-        starting a new sequence (with no prefix).
-
-    Returns:
-      A state object that can be passed to __call__.
-    """
-
-    if not self.use_causal_mask:
-      raise ValueError("Generator must have been trained with a causal mask.")
-
-    # Get relative position bias.
-    rel_position_bias = self.relative_positions(
-        1, self.window_length, offset=self.window_length, bidirectional=False
-    )
-    rel_position_bias = jnp.tile(rel_position_bias, (self.batch_size, 1, 1, 1))
-
-    # Initialize autoregressive storage for (key, value) pairs.
-    # Include space for a prefix of window_length tokens.
-    num_keys = sequence_length + self.window_length
-    stored_shape = (self.batch_size, num_keys, self.num_heads, self.head_size)
-    stored_keys = jnp.zeros(stored_shape, dtype=self.dtype)
-    stored_values = jnp.zeros(stored_shape, dtype=self.dtype)
-
-    recurrent_kvq = None
-    current_index = jnp.array([self.window_length] * self.batch_size)
-
-    decoder_state_dict = {
-        "keys": stored_keys,
-        "values": stored_values,
-        "current_index": current_index,
-        "relative_position_bias": rel_position_bias,
-        "recurrent_kvq": recurrent_kvq,
-    }
-    return DecoderState(decoder_state_dict)
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""A single transformer layer in inference mode.
+
+Modified
+https://github.com/google-research/meliad/blob/main/transformer/transformer_layer.py
+To accommodate sequence packing + kv cache + relative position during test time.
+"""
+
+from typing import Callable, Mapping, NewType, Optional, Tuple
+
+from absl import logging
+import gin
+import jax
+import jax.numpy as jnp
+from meliad_lib.meliad.transformer import attention
+from meliad_lib.meliad.transformer import nn_components
+from meliad_lib.meliad.transformer import position
+from meliad_lib.meliad.transformer import transformer_layer
+
+Array = jnp.ndarray
+DecoderState = NewType("DecoderState", Mapping[str, Array])
+WindowState = Optional[Tuple[attention.KVITuple, Array]]
+
+
+@jax.vmap
+def update_slice_in_dim_1(array: Array, update: Array, idx: Array) -> Array:
+  """Update a stored keys/values slice for different-lengthed seqs in batch."""
+  return jax.lax.dynamic_update_slice_in_dim(array, update, idx, axis=0)
+
+
+def slice_in_dim_1(window_length: int) -> Callable[[Array, Array], Array]:
+  @jax.vmap
+  def fn(array: Array, idx: Array) -> Array:
+    return jax.lax.dynamic_slice_in_dim(array, idx, window_length, axis=0)
+
+  return fn
+
+
+@gin.configurable
+class TransformerLayerGenerate(transformer_layer.TransformerLayer):
+  """Full transformer layer, with attention."""
+
+  def _next_decoder_state(
+      self, decoder_state: DecoderState, keys: Array, values: Array
+  ) -> Tuple[DecoderState, Array, Array]:
+    """Compute the next decoder state, and return keys,values to attend to.
+
+    The keys,values returned from this function are drawn from the prior
+    decoding state, and comprise a full window of local context.
+
+    Args:
+      decoder_state: The current decoder state, initially created using
+        init_decoder_state().
+      keys: The key for the current token, of shape (batch_size, 1, dim)
+      values: The value for the current token of shape (batch_size, 1, dim)
+
+    Returns:
+      (next_decoder_state,
+       window of keys of shape (batch_size, window_length, dim),
+       window of values of shape (batch_size, window_length, dim))
+    """
+
+    assert keys.shape[1] == 1  # single-token autoregressive decoding.
+
+    # Unpack decoder_state
+    stored_keys = decoder_state["keys"]
+    stored_values = decoder_state["values"]
+    curr_index = decoder_state["current_index"]
+
+    # Slice to get window_length-sized chunk of previous keys,values.
+    out_decoder_state = {}
+    curr_win_index = curr_index - self.window_length
+
+    # out_keys = jax.lax.dynamic_slice_in_dim(
+    #     stored_keys, curr_win_index, self.window_length, axis=1)
+    out_keys = slice_in_dim_1(self.window_length)(stored_keys, curr_win_index)
+
+    # out_values = jax.lax.dynamic_slice_in_dim(
+    #     stored_values, curr_win_index, self.window_length, axis=1)
+    out_values = slice_in_dim_1(self.window_length)(
+        stored_values, curr_win_index
+    )
+
+    # Write current keys,values to stored keys, values.
+    # stored_keys = jax.lax.dynamic_update_slice_in_dim(
+    #     stored_keys, keys, curr_index, axis=1)
+    stored_keys = update_slice_in_dim_1(stored_keys, keys, curr_index)
+    # stored_values = jax.lax.dynamic_update_slice_in_dim(
+    #     stored_values, values, curr_index, axis=1)
+    stored_values = update_slice_in_dim_1(stored_values, values, curr_index)
+    curr_index = curr_index + 1
+
+    # Pack a new decoder_state object.
+    out_decoder_state["keys"] = stored_keys
+    out_decoder_state["values"] = stored_values
+    out_decoder_state["current_index"] = curr_index
+    out_decoder_state["relative_position_bias"] = decoder_state[
+        "relative_position_bias"
+    ]
+    out_decoder_state["recurrent_kvq"] = decoder_state["recurrent_kvq"]
+
+    return (DecoderState(out_decoder_state), out_keys, out_values)
+
+  def __call__(
+      self,
+      xs: Array,
+      start_of_sequence: Array,
+      *,
+      importance: Optional[Array] = None,
+      cross_attention_kv: Optional[Tuple[Array, Array]] = None,
+      window_state: Optional[WindowState] = None,
+      decoder_state: Optional[DecoderState] = None,
+  ):
+    """Computes attention over a sequence of inputs.
+
+    Args:
+      xs: input sequence of shape (batch_size, sequence_length, num_hidden)
+      start_of_sequence: An input array of shape (batch_size)  --- The following
+        must be passed by keyword only. ---
+      importance: Array of shape (batch_size, sequence_length). An importance
+        bias for attention.
+      cross_attention_kv: Keys and values from encoder for cross-attention.
+      window_state: State object which contains context from the prior window
+        when using a transformer-XL or sliding window. Initially created with
+        load_window_state().
+      decoder_state: State object for autoregressive decoding, initially created
+        with from init_decoder_state().
+
+    Returns:
+      (ys: outputs of shape (batch_size, sequence_length, num_hidden),
+       importance_score: importance score for the next layer,
+       next_window_state: state to pass to the next window,
+       next_decoder_state: next decoder state for autoregressive decoding,
+       viz_dict: dictionary of visualizations
+      )
+    """
+
+    xs = jnp.asarray(xs, dtype=self.dtype)
+    logging.info("tlayer: recurrent = %r", self.recurrent_attention)
+    logging.info("tlayer: compute_importance = %r", self.compute_importance)
+
+    is_training = self.mode == "train"
+
+    # Compute keys, values and queries.
+    # ---------------------------------
+    logging.info("tlayer: compute keys,values,queries.")
+    (keys, values, queries, queries2) = self.tbase.kvq(xs)
+    attention_scale_factors = self.tbase.attention_scale_factors()
+    (_, sequence_length, num_heads, _) = queries.shape  # (b, k, h, d)
+
+    # Get biases and masks that are shared across windows.
+    # ----------------------------------------------------
+    if decoder_state is not None:
+      logging.info("tlayer: using autoregressive decoder.")
+      # When decoding, prior keys,values are loaded from the decoder state.
+      # Other values are precomputed, and loaded from the decoder state.
+      # The decoder state will be updated with the current token.
+      assert window_state is None
+
+      prev_kvi = None
+      recurrent_state = None  # Use precomputed recurrent_kvq.
+      cross_attention_kv = None
+      rel_position_bias = decoder_state["relative_position_bias"]
+      causal_mask = None
+      dropout_multiplier = None
+
+      # Reuse cached recurrent keys,values for each token.
+      cached_recurrent_kvq = decoder_state["recurrent_kvq"]
+      if cached_recurrent_kvq is not None:
+        assert cross_attention_kv is None
+        cross_attention_kv = (cached_recurrent_kvq[0], cached_recurrent_kvq[1])
+      del cached_recurrent_kvq
+
+      # Get a full window of keys,values and update decoder state.
+      (decoder_state, keys, values) = self._next_decoder_state(
+          decoder_state, keys, values
+      )
+
+      # Each query attends to window_length prior keys.
+      assert keys.shape[1] == self.window_length
+      kq_relative_offset = self.window_length
+
+      if not self.use_long_xl_architecture:
+        kqpos = position.relative_positions(
+            1, self.window_length, offset=0
+        )  # 2D mask
+        current_idx = decoder_state["current_index"]
+
+        # add (batch, heads) dims for kqpos
+        kqpos = jnp.expand_dims(kqpos, axis=(0, 1))
+        kqpos = jnp.tile(kqpos, (1, self.num_heads, 1, 1))
+
+        # add (_, heads, _) dim for current_idx
+        current_idx = jnp.expand_dims(current_idx, axis=(1, 2, 3))
+
+        causal_mask = kqpos > self.window_length * 2 - current_idx
+    else:
+      logging.info("tlayer: windowed attention.")
+      # When training, attention is done using windows or chunks, and prior
+      # context (e.g. keys,values from the previous window) is stored in the
+      # window_state object.
+      (prev_kvi, recurrent_state) = (
+          window_state  # pytype: disable=attribute-error
+      )
+
+      # Get the size of the sliding window for pos bias, dropout, & causal mask.
+      (num_queries, num_keys) = attention.sliding_attention_window_shape(
+          (keys, values, importance),
+          prev_kvi,
+          queries,
+          window_length=self.window_length,
+      )
+      kq_relative_offset = num_keys - num_queries
+
+      # Get the relative position bias.
+      # The bias doesn't depend on the query content, and so can be precomputed.
+      if self.relative_positions is not None:
+        rel_position_bias = self.relative_positions(
+            num_queries, num_keys, bidirectional=False
+        )
+      else:
+        rel_position_bias = None
+
+      # Get causal mask.
+      if self.use_causal_mask:
+        causal_mask = position.causal_mask(
+            num_queries, num_keys, window_length=self.window_length
+        )
+      else:
+        causal_mask = None
+
+      # Apply dropout to the attention matrix.
+      # The mask will be broadcast across batches and windows.
+      if self.attn_dropout_rate > 0.0 and is_training:
+        dropout_rng = self.make_rng("dropout")
+        attn_shape = (self.num_heads, num_queries, num_keys)
+        dropout_multiplier = nn_components.dropout_multiplier_mask(
+            dropout_rng, self.attn_dropout_rate, attn_shape, self.dtype
+        )
+      else:
+        dropout_multiplier = None
+
+    # Load and store values into external memory, if memory is not None.
+    # ------------------------------------------------------------------
+    (mode, _, update_memory) = self._get_cache_name_from_mode(self.mode)
+    external_kv = self._query_external_memory(
+        keys,
+        values,
+        queries,
+        start_of_sequence=start_of_sequence,
+        mode=mode,
+        update_memory=decoder_state is None and update_memory,
+    )
+
+    if (
+        self.memory is not None
+        and self.memory_combine_with_local == "TRAINABLE_WEIGHTED_MEAN"
+    ):
+      external_memory_bias = jnp.asarray(self.memory_bias, dtype=self.dtype)
+      external_memory_bias = jnp.reshape(
+          external_memory_bias, (1, 1, num_heads, 1)
+      )
+      external_memory_bias = jax.nn.sigmoid(external_memory_bias)
+    else:
+      external_memory_bias = None
+
+    # Compute the number of windows.
+    # ------------------------------
+    if sequence_length < self.window_length:
+      num_windows = 1  # Happens with autoregressive decoding.
+    elif sequence_length == self.window_length:
+      num_windows = 1
+      if self.use_long_xl_architecture:
+        assert prev_kvi is not None
+    else:
+      if not self.use_long_xl_architecture:
+        raise ValueError("Can only use sliding window with Transformer XL.")
+      num_windows = sequence_length // self.window_length
+      if (num_windows * self.window_length) != sequence_length:
+        raise ValueError(
+            f"Window length {self.window_length} must be a "
+            + f"multiple of sequence length {sequence_length}"
+        )
+    logging.info("tlayer: num_windows = %d.", num_windows)
+
+    # Define the function to do attention within a single window.
+    # ---------------------------------------------------------
+    def single_window_attention(
+        carry: tuple[Array, Array], inputs_w: tuple[Array, Array]
+    ) -> tuple[tuple[Array, Array], tuple[Array, Array]]:
+      # This function uses the following variables from the outer scope.
+      # They are listed here for clarity.
+      nonlocal rel_position_bias
+      nonlocal causal_mask
+      nonlocal kq_relative_offset
+      nonlocal dropout_multiplier
+      nonlocal attention_scale_factors
+      nonlocal external_memory_bias
+      nonlocal cross_attention_kv  # externally supplied.
+
+      # keys,values,queries over the whole sequence will be split into chunks.
+      # xs_w, kvqi_w, etc. are the chunk for the current window.
+      (prev_kvi_w, rec_state) = carry  # carried from one window to the next.
+      (kvqi_w, external_kv_w) = inputs_w  # inputs to the current window.
+      # (keys_curr_w, values_curr_w, _, _, importance_curr_w) = kvqi_w
+
+      # Concatenate keys,values from the previous window with the current
+      # window to implement sliding window attention.
+      (kvqi_w, next_kvi_w) = attention.concat_kvqi(kvqi_w, prev_kvi_w)
+      (keys_w, values_w, queries_w, queries2_w, importance_w) = kvqi_w
+
+      # Perform recurrent attention within the current window to get the next
+      # recurrent state, and set up cross attention.
+      if rec_state is not None:
+        logging.info("tlayer: recurrent attention.")
+
+        # NOTE -- recurrent states and input tokens are handled separately,
+        # because they have separate learned positional embeddings.  Due to
+        # the way TransformerBase does cross-attention, this means that we use
+        # separate key,value layers for rec_state and tokens_w.
+
+        # Keys, values, queries from recurrent state.
+        logging.info("tlayer: recurrent kvq.")
+        rec_kvq = self.recurrent_tbase.kvq(rec_state)
+        r_scale_factors = self.recurrent_tbase.attention_scale_factors()
+        (r_keys, r_values, r_queries, r_queries2) = rec_kvq
+
+        # Joint attention over both recurrent states and input tokens.
+        logging.info("tlayer: recurrent self-attention.")
+        r_attn_ys = attention.simple_attention(
+            r_keys,
+            r_values,
+            r_queries,
+            None,
+            scale_factor=r_scale_factors[0],
+            dtype=self.dtype,
+        )
+
+        logging.info("tlayer: recurrent cross-attention.")
+        r_cross_attn_ys = attention.simple_attention(
+            keys_w,
+            values_w,
+            r_queries2,
+            importance_w,
+            scale_factor=r_scale_factors[1],
+            dtype=self.dtype,
+        )
+
+        # Recurrent post-attention FFN.
+        logging.info("tlayer: recurrent ffn.")
+        next_rec_state = self.recurrent_tbase.post_attn_ffn(
+            rec_state, r_attn_ys, r_cross_attn_ys
+        )
+
+        # Get keys and values for cross-attention from recurrent state.
+        assert cross_attention_kv is None
+        local_cross_attention_kv = (r_keys, r_values)
+      else:
+        # Get keys and values for cross-attention from external argument.
+        next_rec_state = None
+        local_cross_attention_kv = cross_attention_kv
+
+      # If using RoPE, keys and queries are rotated before self-attention.
+      if self.relative_position_type == "rotary":
+        logging.info(
+            "Using rotary position encodings (RoPE), offset = %d",
+            kq_relative_offset,
+        )
+        (keys_w, queries_w) = position.rotate_kq(
+            keys_w, queries_w, max_wavelength=10_000, offset=kq_relative_offset
+        )
+
+      # Self-attention over input tokens.
+      logging.info("tlayer: self-attention.")
+      attn_ys_w = attention.simple_attention(
+          keys_w,
+          values_w,
+          queries_w,
+          importance_w,
+          relative_position_bias=rel_position_bias,
+          scale_factor=attention_scale_factors[0],
+          causal_mask=causal_mask,
+          dropout_multiplier=dropout_multiplier,
+          dtype=self.dtype,
+      )
+
+      # Attention over external memory.
+      if external_kv_w is not None:
+        (external_keys_w, external_values_w) = external_kv_w
+        y_ext = attention.external_attention(
+            external_keys_w,
+            external_values_w,
+            queries_w,
+            scale_factor=attention_scale_factors[0],
+        )
+        if external_memory_bias is not None:
+          ebias = external_memory_bias
+          attn_ys_w = (attn_ys_w * (1 - ebias)) + (y_ext * ebias)
+        elif self.memory_combine_with_local == "ADD":
+          attn_ys_w += y_ext
+        elif self.memory_combine_with_local == "STOP_FORWARD":
+          attn_ys_w = y_ext + (attn_ys_w - jax.lax.stop_gradient(attn_ys_w))
+        else:
+          raise ValueError(
+              f"Unexpected setting: {self.memory_combine_with_local = }"
+          )
+
+      # Cross attention from input tokens to encoder or recurrent state.
+      if local_cross_attention_kv is not None:
+        logging.info("tlayer: cross-attention.")
+        (c_keys, c_values) = local_cross_attention_kv
+
+        # Cross-attention using queries2.
+        cross_attn_ys_w = attention.simple_attention(
+            c_keys,
+            c_values,
+            queries2_w,
+            None,
+            scale_factor=attention_scale_factors[1],
+            dtype=self.dtype,
+        )
+      else:
+        cross_attn_ys_w = None
+
+      # End function single_window_attention(...)
+      return ((next_kvi_w, next_rec_state), (attn_ys_w, cross_attn_ys_w))
+
+    # Initialize recurrent_tbase before calling jax.lax.scan.
+    # Otherwise flax will throw a tantrum.
+    if (
+        self.recurrent_attention
+        and 0 <= self.max_unrolled_windows
+        and self.max_unrolled_windows < num_windows
+    ):
+      logging.info("tlayer: force initialization of recurrent_tbase.")
+      self.recurrent_tbase.force_init(recurrent_state)
+
+    # Perform sliding window attention over all keys,values,queries.
+    # --------------------------------------------------------------
+    initial_carry = (prev_kvi, recurrent_state)  # window state.
+    kvqi = (keys, values, queries, queries2, importance)
+    attn_inputs = (kvqi, external_kv)
+    (next_carry, attn_outputs) = attention.split_and_scan(
+        single_window_attention,
+        initial_carry,
+        attn_inputs,
+        sections=num_windows,
+        axis=1,
+        max_unrolled_windows=self.max_unrolled_windows,
+    )
+    (attn_ys, cross_attn_ys) = attn_outputs
+
+    logging.info("tlayer: End windows.")
+
+    # Post-attention MLP, resnet, and FFN.
+    # ------------------------------------
+    logging.info("tlayer: final FFN.")
+    ys = self.tbase.post_attn_ffn(xs, attn_ys, cross_attn_ys)
+
+    # Compute importance scores for each token if requested.
+    if self.compute_importance:
+      (batch_size, sequence_length, _) = ys.shape
+      importance_score = self.importance_layer(ys)
+      importance_score = importance_score.reshape((batch_size, sequence_length))
+    else:
+      importance_score = None
+
+    next_window_state = next_carry if window_state is not None else None
+    viz_dict = {}  # Visualizations, not currently enabled.
+    return (ys, importance_score, next_window_state, decoder_state, viz_dict)
+
+  def init_decoder_state_vanilla(
+      self, sequence_length: int, start_of_sequence: Array
+  ) -> DecoderState:
+    """Initialize decoder state for autoregressive generation.
+
+    Args:
+      sequence_length: The maximum length of the sequence to generate.
+      start_of_sequence: Array of boolean of shape (batch_size,) True if
+        starting a new sequence (with no prefix).
+
+    Returns:
+      A state object that can be passed to __call__.
+    """
+
+    if not self.use_causal_mask:
+      raise ValueError("Generator must have been trained with a causal mask.")
+
+    # Get relative position bias.
+    rel_position_bias = self.relative_positions(
+        1, self.window_length, offset=self.window_length, bidirectional=False
+    )
+    rel_position_bias = jnp.tile(rel_position_bias, (self.batch_size, 1, 1, 1))
+
+    # Initialize autoregressive storage for (key, value) pairs.
+    # Include space for a prefix of window_length tokens.
+    num_keys = sequence_length + self.window_length
+    stored_shape = (self.batch_size, num_keys, self.num_heads, self.head_size)
+    stored_keys = jnp.zeros(stored_shape, dtype=self.dtype)
+    stored_values = jnp.zeros(stored_shape, dtype=self.dtype)
+
+    recurrent_kvq = None
+    current_index = jnp.array([self.window_length] * self.batch_size)
+
+    decoder_state_dict = {
+        "keys": stored_keys,
+        "values": stored_values,
+        "current_index": current_index,
+        "relative_position_bias": rel_position_bias,
+        "recurrent_kvq": recurrent_kvq,
+    }
+    return DecoderState(decoder_state_dict)