make_sample.py

#!/usr/bin/env python3
"""
Generate noisy point clouds that mimic 3D-generation output, so the refiner
has something realistic to chew on.

    python make_sample.py torus  -o samples/torus.ply  -n 40000 --noise 0.01
    python make_sample.py sphere -o samples/sphere.ply
    python make_sample.py blob   -o samples/blob.ply

These come out as raw, normal-less, noisy point clouds -- exactly the kind of
mess a generator hands you.
"""
import argparse
import numpy as np


def sphere(n, rng):
    v = rng.normal(size=(n, 3))
    v /= np.linalg.norm(v, axis=1, keepdims=True)
    return v


def torus(n, rng, R=1.0, r=0.38):
    u = rng.uniform(0, 2 * np.pi, n)
    v = rng.uniform(0, 2 * np.pi, n)
    x = (R + r * np.cos(v)) * np.cos(u)
    y = (R + r * np.cos(v)) * np.sin(u)
    z = r * np.sin(v)
    return np.stack([x, y, z], axis=1)


def blob(n, rng):
    """Lumpy sphere -- varying curvature, the interesting case for quad flow."""
    p = sphere(n, rng)
    bump = (1.0
            + 0.25 * np.sin(3 * p[:, 0] * np.pi)
            + 0.18 * np.cos(4 * p[:, 1] * np.pi)
            + 0.12 * np.sin(5 * p[:, 2] * np.pi))
    return p * bump[:, None]


SHAPES = {"sphere": sphere, "torus": torus, "blob": blob}


def write_ply(path, pts):
    pts = np.asarray(pts, dtype="float32")
    with open(path, "w", newline="\n") as f:
        f.write("ply\nformat ascii 1.0\n")
        f.write(f"element vertex {len(pts)}\n")
        f.write("property float x\nproperty float y\nproperty float z\n")
        f.write("end_header\n")
        for x, y, z in pts:
            f.write(f"{x:.6f} {y:.6f} {z:.6f}\n")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("shape", choices=list(SHAPES))
    ap.add_argument("-o", "--output", required=True)
    ap.add_argument("-n", "--count", type=int, default=40000)
    ap.add_argument("--noise", type=float, default=0.012,
                    help="gaussian noise as fraction of unit scale")
    ap.add_argument("--seed", type=int, default=7)
    a = ap.parse_args()

    rng = np.random.default_rng(a.seed)
    pts = SHAPES[a.shape](a.count, rng)
    pts = pts + rng.normal(scale=a.noise, size=pts.shape)
    write_ply(a.output, pts)
    print(f"wrote {len(pts):,} points -> {a.output}")


if __name__ == "__main__":
    main()