Upload 2 files

ag4masses/alphageometry/alphageometry.py

@@ -33,6 +33,8 @@ 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(
@@ -199,12 +201,11 @@ def write_solution(g: gh.Graph, p: pr.Problem, out_file: str) -> None:
     rule_name = r2name.get(con.rule_name, '')
     nl = nl.replace('\u21d2', f'{rule_name}\u21d2 ')
     solution += '{:03}. '.format(i + 1) + nl + '\n'
-  logging.info(solution)
+  # logging.info(solution)
   if out_file:
     with open(out_file, 'w') as f:
       f.write(solution)
-    logging.info('Solution written to %s.', out_file)
-
+    # logging.info('Solution written to %s.', out_file)
 
 def get_lm(ckpt_init: str, vocab_path: str) -> lm.LanguageModelInference:
   lm.parse_gin_configuration(
@@ -233,7 +234,6 @@ 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],
@@ -598,7 +598,7 @@ def bqsearch(i_nd, srch_inputs, out_file) -> tuple[int, bool, list]: # ( iNode,
     return (i_nd, False, ret)
 
 def run_alphageometry(
-    #XX model: lm.LanguageModelInference,
+    # model: lm.LanguageModelInference,
     p: pr.Problem,
     search_depth: int,
     beam_size: int,
@@ -739,9 +739,9 @@ def main(_):
     run_ddar(g, this_problem, _OUT_FILE.value)
 
   elif _MODE.value == 'alphageometry':
-    #XX model = get_lm(_CKPT_PATH.value, _VOCAB_PATH.value)
+    model = get_lm(_CKPT_PATH.value, _VOCAB_PATH.value)
     run_alphageometry(
-        #XX model,
+        model,
         this_problem,
         _SEARCH_DEPTH.value,
         _BEAM_SIZE.value,

ag4masses/alphageometry/numericals.py

@@ -29,6 +29,10 @@ from numpy.random import uniform as unif  # pylint: disable=g-importing-member
 import graph as gh
 from collections import defaultdict
 from itertools import combinations
+import numpy as np
+import matplotlib.patches
+import matplotlib.pyplot as plt
+from itertools import combinations
 
 matplotlib.use('TkAgg')
 
@@ -1060,7 +1064,7 @@ def _draw_line(
 
 
 def draw_line(
-    ax: matplotlib.axes.Axes, line: Line, color: Any = 'white'
+    ax: matplotlib.axes.Axes, line: Line, color: Any = 'white', draw: bool = True
 ) -> tuple[Point, Point]:
   """Draw a line."""
   points = line.neighbors(gm.Point)
@@ -1080,8 +1084,11 @@ def draw_line(
       pmax = p, v
 
   p1, p2 = pmin[0], pmax[0]
-  _draw_line(ax, p1, p2, color=color)
-  return p1, p2
+  if draw:
+    _draw_line(ax, p1, p2, color=color)
+    return p1, p2
+  else:
+    return p1, p2
 
 
 def _draw_circle(
@@ -1315,10 +1322,6 @@ def highlight(
 ) -> None:
   """Draw highlights."""
   args = list(map(lambda x: x.num if isinstance(x, gm.Point) else x, args))
-
-  if name == 'cyclic':
-    a, b, c, d = args
-    _draw_circle(ax, Circle(p1=a, p2=b, p3=c), color=color1, lw=2.0)
   if name == 'coll':
     a, b, c = args
     a, b = max(a, b, c), min(a, b, c)
@@ -1407,6 +1410,90 @@ def find_pairs_with_same_distance(line_lengths):
     
     return result
 
+def calculate_angle(p1, p2, p3, p4):
+    """Calculates the angle between two lines formed by points (p1, p2) and (p3, p4) in degrees."""
+    # Determine the common point
+    if p2 == p3:
+        common = p2
+        other_points = (p1, p4)
+        v1 = np.array([p1.x - p2.x, p1.y - p2.y])
+        v2 = np.array([p4.x - p2.x, p4.y - p2.y])
+    elif p1 == p4:
+        common = p1
+        other_points = (p2, p3)
+        v1 = np.array([p2.x - p1.x, p2.y - p1.y])
+        v2 = np.array([p3.x - p1.x, p3.y - p1.y])
+    elif p1 == p3:
+        common = p1
+        other_points = (p2, p4)
+        v1 = np.array([p2.x - p1.x, p2.y - p1.y])
+        v2 = np.array([p4.x - p1.x, p4.y - p1.y])
+    elif p2 == p4:
+        common = p2
+        other_points = (p1, p3)
+        v1 = np.array([p1.x - p2.x, p1.y - p2.y])
+        v2 = np.array([p3.x - p2.x, p3.y - p2.y])
+    else:
+        return None, None, None  # No shared point, angle cannot be calculated
+
+    # Calculate the angle
+    cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
+    cos_angle = np.clip(cos_angle, -1, 1)  # Ensure valid range for acos
+    angle_rad = np.arccos(cos_angle)
+    return common, other_points, round(np.degrees(angle_rad), 5)
+
+def highlight_angle2(ax, origin, p1, p2, radius, color):
+    """Highlights the angle formed by two vectors meeting at 'origin'."""
+    # Calculate angles of vectors
+    angle1 = np.arctan2(p1.y - origin.y, p1.x - origin.x)
+    angle2 = np.arctan2(p2.y - origin.y, p2.x - origin.x)
+
+    # Convert to degrees and ensure the smaller angle is first
+    angle1_deg, angle2_deg = sorted(np.degrees([angle1, angle2]))
+    if angle2_deg - angle1_deg > 180:
+        angle1_deg, angle2_deg = angle2_deg, angle1_deg
+    
+    # Draw the wedge
+    wedge = matplotlib.patches.Wedge(
+        center=(origin.x, origin.y),
+        r=radius,
+        theta1=angle1_deg,
+        theta2=angle2_deg,
+        color=color,
+        alpha=0.5
+    )
+    ax.add_patch(wedge)
+    # print("Angle highlighted with color:", color)
+
+def search_in_dict(num, my_dict):
+    for key in my_dict.keys():
+        if round(key, 3) == round(num, 3):
+            return True
+    return False
+
+def highlight_same_angle(ax, lines, color_list):
+    """Highlights angles formed at shared points by pairs of lines."""
+    lines_list = [(draw_line(ax, l, draw=False)) for l in lines]  # Extract points for all lines
+    angle_color_radius = {}
+
+    for line1, line2 in combinations(lines_list, 2):
+        # Calculate the angle
+        common_point, other_points, angle = calculate_angle(*line1, *line2)
+        if angle is None or angle > 90:
+            continue  # Skip invalid or small angles
+
+        if search_in_dict(angle, angle_color_radius) == False:
+          # Assign color and radius for this unique angle
+          color = color_list[len(angle_color_radius) % len(color_list)]
+          radius = 0.1 + len(angle_color_radius) * 0.05
+          angle_color_radius[round(angle,3)] = (color, radius)
+          # print(type(angle))
+        else:
+          color, radius = angle_color_radius[round(angle, 3)]
+          # print(type(angle))
+
+        # Highlight the angle
+        highlight_angle2(ax, common_point, *other_points, radius, color)
 def _draw(
     ax: matplotlib.axes.Axes,
     points: list[gm.Point],
@@ -1471,6 +1558,7 @@ def _draw(
 
     # Call the highlight function with the determined color
     highlight(ax, 'cong', [p1, p2, p3, p4], lcolor, color, color)
+    highlight_same_angle(ax, lines, color_list=colors_highlight)
   if equals:
     for i, segs in enumerate(equals['segments']):
       color = colors[i % len(colors)]