FPO Mesh Gen

mapgen.py

from pathlib import Path
import numpy as np
import numba as nb
import scipy
import pickle
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from collections import Counter


def pickle_load(filename):
    """Loads any data from file via unpickling. Use at your own risk!"""
    with open(filename, 'rb') as f:
        data = pickle.load(f)
    return data


def pickle_save(filename, data):
    """Saves any data to file via pickling."""
    with open(filename, 'wb') as f:
        pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)


@nb.njit(nb.float64(nb.float64[:, :]))
def polygon_area(points):
    N = len(points)
    area = np.float64(0)
    for i in range(N):
        xi = points[i][0]
        ya = points[i - 1][1]
        yb = points[(i + 1) % N][1]
        area += xi * (yb - ya)
    return area / 2


@nb.njit(nb.float64[:](nb.float64[:, :]))
def polygon_centroid(points):
    N = len(points)
    area = np.float64(0)
    centroid = np.array([0, 0], dtype=np.float64)
    for i in range(N):
        xi, yi = points[i - 1]
        xj, yj = points[i]
        p = xi * yj - xj * yi
        area += p
        centroid[0] += (xi + xj) * p
        centroid[1] += (yi + yj) * p

    centroid /= 3 * area
    return centroid


@nb.njit(nb.float64(nb.float64[:], nb.float64[:], nb.float64[:]))
def distance_to_line_segment(p0, p1, p2):
    """Distance of p0 to line segment p1-p2."""
    q0 = p0 - p1
    q2 = p2 - p1
    t = np.dot(q2, q0) / np.dot(q2, q2)

    if t <= 0:
        return np.linalg.norm(q0)
    elif t >= 1:
        return np.linalg.norm(p0 - p2)
    else:
        return np.sqrt(np.dot(q0, q0) - (t**2) * np.dot(q2, q2))


@nb.njit(nb.float64(nb.float64[:, :], nb.float64[:]))
def distance_to_polygon(polygon, point):
    N = len(polygon)
    res = np.inf

    for i in range(N):
        a = polygon[i - 1]
        b = polygon[i]

        res = min(res, distance_to_line_segment(point, a, b))

    return res


@nb.njit(nb.types.Tuple((nb.float64, nb.float64[:]))(nb.float64[:, :]))
def circumcircle(p):
    """Returns radius squared and circumcenter."""
    Bx = p[1, 0] - p[0, 0]
    By = p[1, 1] - p[0, 1]
    Cx = p[2, 0] - p[0, 0]
    Cy = p[2, 1] - p[0, 1]

    D = 2 * (Bx * Cy - By * Cx)
    BB = Bx**2 + By**2
    CC = Cx**2 + Cy**2
    Ux = (Cy * BB - By * CC) / D
    Uy = (Bx * CC - Cx * BB) / D
    return Ux**2 + Uy**2, np.array([Ux, Uy]) + p[0]


@nb.njit(nb.float64[:](nb.float64[:], nb.float64[:], nb.float64[:], nb.float64))
def circumcenter_centroid_point(a, b, c, k):
    """Returns a point in between circumcenter (k=0) and centroid (k=1)."""
    Bx, By = b - a
    Cx, Cy = c - a
    D = 2 * (Bx * Cy - By * Cx)
    BB = Bx**2 + By**2
    CC = Cx**2 + Cy**2
    Ux = (Cy * BB - By * CC) / D
    Uy = (Bx * CC - Cx * BB) / D

    circumcenter = np.array([Ux, Uy])
    centroid = np.array([Bx + Cx, By + Cy]) / 3
    return k * centroid + (1 - k) * circumcenter + a


@nb.njit(nb.bool(nb.float64[:], nb.float64[:], nb.float64[:]))
def is_acute(a, b, c):
    """Returns True if triangle is acute."""
    A = (b[0] - c[0]) ** 2 + (b[1] - c[1]) ** 2
    B = (a[0] - c[0]) ** 2 + (a[1] - c[1]) ** 2
    C = (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2

    if A > C:
        A, C = C, A
    if B > C:
        B, C = C, B

    return C < A + B


@nb.njit(nb.float64(nb.float64[:], nb.float64[:], nb.float64[:]))
def largest_angle(a, b, c):
    """Returns largest angle in triangle."""
    A = (b[0] - c[0]) ** 2 + (b[1] - c[1]) ** 2
    B = (a[0] - c[0]) ** 2 + (a[1] - c[1]) ** 2
    C = (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2

    if A > C:
        A, C = C, A
    if B > C:
        B, C = C, B

    return np.arccos((A + B - C) / 2 / np.sqrt(A * B))


@nb.njit(nb.bool(nb.float64[:]))
def in_torus(p):
    return (0 <= p[0] < 1) and (0 <= p[1] < 1)


@nb.njit(nb.bool[:](nb.float64[:, :]))
def points_in_torus(points):
    """Return True if points are in primary torus."""
    N = len(points)
    res = np.zeros(N, dtype=np.bool)

    for i in range(N):
        x = points[i, 0]
        y = points[i, 1]
        res[i] = ((0.0 <= x < 1.0) & (0.0 <= y < 1.0))

    return res


@nb.njit(nb.float64(nb.float64[:], nb.float64[:]))
def squared_dist_torus(a, b):
    x = np.abs(a[0] - b[0])
    y = np.abs(a[1] - b[1])
    if x > 0.5:
        x = 1.0 - x
    if y > 0.5:
        y = 1.0 - y
    return x * x + y * y


@nb.njit(nb.types.UniTuple(nb.float64, 2)(nb.float64[:, :]), parallel=True)
def min_dist_statistics(points):
    N = len(points)
    D_max = np.sqrt(2 / np.sqrt(3) / N)

    dx = np.ones(N) * np.inf
    for i in nb.prange(N):
        for j in range(N):
            if i != j:
                dist = np.sqrt(squared_dist_torus(points[i], points[j]))
                if dist < dx[i]:
                    dx[i] = dist

    return dx.min() / D_max, dx.mean() / D_max


@nb.njit(nb.float64(nb.float64[:], nb.float64[:]))
def cross_product(v, w):
    return v[0] * w[1] - v[1] * w[0]


@nb.njit(nb.float64(nb.float64[:], nb.float64[:]))
def angle_between(a, b):
    c = a[0] * b[0] + a[1] * b[1]
    n = (a[0] ** 2 + a[1] ** 2) * (b[0] ** 2 + b[1] ** 2)
    return np.arccos(np.abs(c) / np.sqrt(n))


@nb.njit(
    nb.types.UniTuple(nb.float64, 2)(
        nb.float64[:], nb.float64[:], nb.float64[:], nb.float64[:]
    )
)
def line_intersection(p, r, q, s):
    qp = q - p
    rs = cross_product(r, s)

    t = cross_product(qp, s) / rs
    u = cross_product(qp, r) / rs
    return t, u


def get_neighbors(k, tri):
    indptr, indices = tri.vertex_neighbor_vertices
    return indices[indptr[k] : indptr[k + 1]]


def degree(k, tri):
    indptr, _ = tri.vertex_neighbor_vertices
    return int(indptr[k + 1] - indptr[k])


def generate_offsets(num_rings=0):
    res = [np.array([0, 0])]

    for i in range(num_rings):
        res.append(res[-1] + np.array([1, 0]))

        for n, a in (
            (1 + 2*i, np.array([0, -1])),
            (2 + 2*i, np.array([-1, 0])),
            (2 + 2*i, np.array([0, 1])),
            (2 + 2*i, np.array([1, 0])),
        ):
            for _ in range(n):
                res.append(res[-1] + a)

    return np.stack(res)


def periodic_tiling(points, num_rings=1, clip=False):
    N = len(points)

    offsets = generate_offsets(num_rings)
    res = np.zeros_like(points, shape=(len(offsets) * N, 2))

    for i, offset in enumerate(offsets):
        res[i*N:(i+1)*N] = points + offset

    if clip:
        e = 4.0 / np.sqrt(N)

        filt = ((-e <= res) & (res < 1 + e)).all(1)
        res = res[filt]

    return res


def triangulate_periodic(points, clip=False):
    return scipy.spatial.Delaunay(periodic_tiling(points, clip=clip))


def torus_plot_tri(ax, tri, N, clip=None):
    for i in range(-1, 3):
        ax.axvline(i, color="0.5", ls="--", lw=1, zorder=2)
        ax.axhline(i, color="0.5", ls="--", lw=1, zorder=2)

    ax.add_patch(
        plt.Rectangle((0, 0), 1, 1, fill=False, ec="0.3", ls="-", lw=1, zorder=3)
    )
    if clip is None:
        clip = 4.0 / np.sqrt(N)
    ax.set_xlim(-clip, 1 + clip)
    ax.set_ylim(-clip, 1 + clip)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax.plot(*tri.points.T, "C3.", ms=7, zorder=5)
    ax.plot(*tri.points[:N].T, "o", ms=8, mfc="None", mec="C3", zorder=5)
    ax.triplot(
        *tri.points.T, tri.simplices, color="darkgreen", lw=1, zorder=4, alpha=0.8
    )


def get_polygon_patches(mesh, k):
    offsets = np.array(
        [[-1, 0], [-1, 1], [0, 1], [1, 1], [1, 0], [1, -1], [0, -1], [-1, -1]]
    )

    cells = mesh.compute_cells(k)
    min_radius = min(distance_to_polygon(p, x) for p, x in zip(cells, mesh.points))
    circles = [plt.Circle(x, min_radius) for x in mesh.points]

    patches = [plt.Polygon(p) for p in cells]
    other_patches = [plt.Polygon(p + o) for p in cells for o in offsets]

    areas = np.array([polygon_area(p) for p in cells])
    other_areas = np.array([polygon_area(p) for p in cells for o in offsets])

    return patches, areas, other_patches, other_areas, circles


def torus_plot_cells(mesh, k, cmin=0.00, cmax=0.02):
    p_kw = dict(ls="-", ec="C4", lw=1, cmap="viridis")
    patches, areas, other_patches, other_areas, circles = get_polygon_patches(mesh, k)

    pc = PatchCollection(patches, alpha=0.6, **p_kw)
    pc.set_array(areas)
    pc.set_clim(cmin, cmax)

    opc = PatchCollection(other_patches, alpha=0.2, **p_kw)
    opc.set_array(other_areas)
    opc.set_clim(cmin, cmax)

    cpc = PatchCollection(circles, alpha=0.5, fc="None", ec="k", lw=1, ls=":")

    return pc, opc, cpc


class BlueNoiseSamplingFPO:
    def __init__(self, num_points, num_passes, seed, cache=True):
        self.num_points = int(num_points)
        self.num_passes = int(num_passes)
        self.seed = int(seed)
        assert self.num_passes >= 10

        cache_folder = Path("cached")
        cache_folder.mkdir(exist_ok=True)

        _file_name = f"{self.num_points}_points_fpo_{self.seed}.points"
        self._cache = Path("cached") / _file_name

        if self._cache.exists():
            self.point_dict = pickle_load(self._cache)
            if self.num_passes in self.point_dict:
                self.points = self.point_dict[self.num_passes]
                return

        self.generate_points()
        if cache:
            pickle_save(self._cache, self.point_dict)

    def generate_points(self):
        self.RAND = np.random.default_rng(seed=self.seed)
        self.points = self.RAND.random((self.num_points, 2))
        self.point_dict = dict()

        clip = False
        for i in range(1, self.num_passes + 1):
            self.global_pass(clip)
            clip = True

            if (i + 1) % 10 == 0:
                self.point_dict[i + 1] = self.points.copy()

        if self.num_passes not in self.point_dict:
            self.point_dict[self.num_passes] = self.points.copy()

    def global_pass(self, clip=False):
        self.RAND.shuffle(self.points)

        for i in range(self.num_points):
            point = self.points[i]

            temp_points = np.delete(self.points, i, axis=0)
            r_max = min(squared_dist_torus(point, p) for p in temp_points)

            tri = triangulate_periodic(temp_points, clip=clip)
            filt = points_in_torus(tri.points[tri.simplices].mean(1))

            for p in tri.points[tri.simplices[filt]]:
                r, c = circumcircle(p)
                if r > r_max:
                    r_max, point = r, c

            self.points[i] = np.mod(point, 1)


class PeriodicTriMesh:
    def __init__(self, points, num_rings=1, recenter=True):
        self.N = len(points)
        self.points = points
        if recenter:
            self.recenter_points()

        self.tiled_points = periodic_tiling(self.points, num_rings)
        self.tri = scipy.spatial.Delaunay(self.tiled_points)

        self.edges = self.compute_edges()
        self.edge_lengths = np.array(
            [
                np.linalg.norm(self.tri.points[b] - self.tri.points[a])
                for a, b in self.edges
            ]
        )

        self.cell_triangles = [self.get_triangles_around(i) for i in range(self.N)]

    def recenter_points(self):
        x = np.sort(self.points[:, 0])
        x = np.append(x, 1 + x[0])
        i = np.argmax(np.diff(x))
        x_offset = 1 - (x[i] + x[i+1])/2

        y = np.sort(self.points[:, 1])
        y = np.append(y, 1 + y[0])
        i = np.argmax(np.diff(y))
        y_offset = 1 - (y[i] + y[i+1])/2

        self.points = (self.points + np.array([x_offset, y_offset])) % 1

    def compute_edges(self):
        edges = dict()
        for i, s in enumerate(self.tri.simplices):
            a, b, c = sorted(s)

            for x, y in [(a, b), (a, c), (b, c)]:
                if in_torus((self.tri.points[x] + self.tri.points[y]) / 2):
                    edges[(x, y)] = edges.get((x, y), []) + [i]

        return edges

    def get_triangles_around(self, i):
        """Get triangles around vertex i in CCW order."""

        def ccw_neighbor(i, neighbors):
            for j, t in enumerate(neighbors):
                if (t == -1) or (i not in self.tri.simplices[t]):
                    return neighbors[(j + 1) % 3]

        triangles = [self.tri.vertex_to_simplex[i]]
        while True:
            neighbors = self.tri.neighbors[triangles[-1]]
            next_t = ccw_neighbor(i, neighbors)

            if next_t == triangles[0]:
                break
            elif next_t not in triangles:
                triangles.append(next_t)
            else:
                raise ValueError("Could not find polygon around vertex.")

        return np.array(triangles)

    def compute_cells(self, k):
        return [
            np.stack(
                [
                    circumcenter_centroid_point(*self.tri.points[t], k)
                    for t in self.tri.simplices[triangles]
                ]
            )
            for triangles in self.cell_triangles
        ]

    def edge_intersections(self, k):
        intersections, angles = [], []
        for (i, j), (a, b) in self.edges.items():
            p = self.tri.points[i]
            r = self.tri.points[j] - p
            ta = self.tri.points[self.tri.simplices[a]]
            tb = self.tri.points[self.tri.simplices[b]]

            q = circumcenter_centroid_point(*ta, k)
            s = circumcenter_centroid_point(*tb, k) - q

            t, u = line_intersection(p, r, q, s)

            intersections.append([t, u])
            angles.append(angle_between(r, s))

        return np.array(intersections), np.array(angles)

    def degree_counter(self):
        return Counter(degree(i, self.tri) for i in range(self.N))

    def min_dist_metric(self):
        return self.edge_lengths.min() / np.sqrt(2 / np.sqrt(3) / self.N)

    def smallest_inscribed_circle(self, k):
        cells = self.compute_cells(k)
        return min(distance_to_polygon(p, x) for p, x in zip(cells, self.points))

    def normalized_distance_metric(self):
        dijk = DijkstraPath(self)

        mid = np.array([0.5, 0.5])
        center_mask = np.linalg.norm(self.points - mid, axis=1) < 0.25

        normed_dists = []
        for i in np.arange(self.N)[center_mask]:
            dijk.compute(i)
            direct = dijk.direct_distance()
            mask = (0.25 <= direct) & (direct <= 0.5)
            normed_dists.append(dijk.dist[mask] / direct[mask])

        return np.concatenate(normed_dists).mean()


class DijkstraPath:
    def __init__(self, mesh, source=0):
        self.mesh = mesh
        self.npoints = mesh.tri.npoints
        self.compute(source)

    def compute(self, source):
        self.source = source
        self.dist = np.ones(self.npoints) * np.inf
        self.prev = [None] * self.npoints

        self.dist[self.source] = 0.0
        Q = set(range(self.npoints))
        while Q:
            u = min(Q, key=lambda x: self.dist[x])
            Q.remove(u)

            for v in get_neighbors(u, self.mesh.tri):
                uv = np.linalg.norm(self.mesh.tri.points[u] - self.mesh.tri.points[v])
                alt = self.dist[u] + uv
                if alt < self.dist[v]:
                    self.dist[v] = alt
                    self.prev[v] = u

    def path_to(self, target):
        path = [target]
        while (u := self.prev[path[-1]]) is not None:
            path.append(u)
        return np.stack([self.mesh.tri.points[i] for i in path])

    def direct_distance(self):
        points = self.mesh.tri.points
        direct = np.linalg.norm(points - points[self.source], axis=1)
        return direct


def print_mesh_metrics(mesh, ks=np.linspace(0.0, 1.0, 11)):
    min_dist = mesh.min_dist_metric()
    path_len = mesh.normalized_distance_metric()

    print(f"N = {mesh.N} (Min Dist = {min_dist:.3f})", end="")
    print(f" | Path Length Metric = {path_len:.3f}")

    deg = mesh.degree_counter()
    small_deg = sum(deg[i] for i in deg if i < 5)
    large_deg = sum(deg[i] for i in deg if i > 7)
    print("Degrees (<5, 5, 6, 7, >7) = ", end="")
    print(f"({small_deg}, {deg[5]}, {deg[6]}, {deg[7]}, {large_deg})\n")

    print("   k  |         Edge Intersection         | Radius |")
    print("------|-----------------------------------|--------|")

    for k in ks:
        r = mesh.N * np.sqrt(12) * mesh.smallest_inscribed_circle(k)**2
        intersections, angles = mesh.edge_intersections(k)
        ang = angles.min() * 180.0 / np.pi

        x = 2 * np.abs(intersections - 0.5)
        (u_2s, u_max), (t_2s, t_max) = np.quantile(x, [0.99, 1.0], axis=0).T
        int_str = f"{u_2s:.2f} - {t_2s:.2f} | {u_max:.2f} - {t_max:5.2f}"
        print(f" {k:.2f} | {int_str} / {ang:.1f} | {r:.4f} |")


def save_periodic_tile(filename, mesh):
    p = mesh.points.astype(np.float32)
    num_points = np.int32(len(p))

    t = np.stack([s for s in mesh.tri.simplices if min(s) < 100]).astype(np.int32)
    num_tris = np.int32(len(t))

    with open(filename, 'wb') as f:
        f.write(num_points.tobytes() + num_tris.tobytes())
        f.write(p.tobytes())
        f.write(t.tobytes())

    print(f"Saved to {filename}")


if __name__ == "__main__":
    N, P, S = 100, 50, 171

    fpo = BlueNoiseSamplingFPO(num_points=N, num_passes=P, seed=S)
    mesh = PeriodicTriMesh(fpo.points)
    save_periodic_tile("base.periodic", mesh)