Failing to understand the planar unit-distance problem with a toy demonstration

unit_distance_difference.py

"""Toy demonstration of the planar unit-distance idea.

This does NOT reproduce OpenAI's construction. It shows the mechanism:
if an arithmetic system gives you more same-length difference vectors, then after
rescaling those vectors to length 1, the point set has more unit-distance pairs.
"""

from math import sqrt


def square_norm(a, b):
    # Gaussian integers a + bi: the square-grid length squared.
    return a * a + b * b


def triangular_norm(a, b):
    # Toy richer arithmetic: triangular/Eisenstein lattice.
    # Vector = a*(1,0) + b*(1/2, sqrt(3)/2)
    # length^2 = a^2 + a*b + b^2
    return a * a + a * b + b * b


def vectors_with_norm(norm, target, search=30):
    return sorted(
        (a, b)
        for a in range(-search, search + 1)
        for b in range(-search, search + 1)
        if (a, b) != (0, 0) and norm(a, b) == target
    )


def count_edges(width, height, vectors):
    points = {(x, y) for x in range(width) for y in range(height)}
    edges = set()

    for x, y in points:
        for dx, dy in vectors:
            q = (x + dx, y + dy)
            if q in points:
                edges.add(tuple(sorted(((x, y), q))))

    return len(edges)


def explain(label, norm, target, width=25, height=25):
    vectors = vectors_with_norm(norm, target)
    scale = sqrt(target)

    print("\n" + label)
    print("squared length before rescaling:", target)
    print("rescale all coordinates by:", f"1/{scale:.3f}")
    print("same-length directions:", len(vectors))
    print("sample directions:", vectors[:12])
    print(
        f"unit-distance edges in a {width}x{height} finite patch:",
        count_edges(width, height, vectors),
    )


if __name__ == "__main__":
    explain(
        "1. Plain square grid: only immediate neighbors count",
        square_norm,
        1,
    )
    explain(
        "2. Clever square/Gaussian grid: hidden diagonal jumps also count",
        square_norm,
        65,
    )
    explain(
        "3. Toy richer arithmetic: more same-length directions",
        triangular_norm,
        91,
    )