331 lines
11 KiB
Python
331 lines
11 KiB
Python
import argparse, numpy as np, os
|
|
from stl.stl import ASCII
|
|
from stl import mesh
|
|
from squish import DomainParams, Simulation, ordered
|
|
from squish.common import OUTPUT_DIR
|
|
from scipy.spatial import Voronoi
|
|
|
|
SUBDIVISION_AMOUNT = 8
|
|
|
|
|
|
def not_in_order(i, j):
|
|
if i < j:
|
|
return j - i != 1
|
|
else:
|
|
return i - j == 1
|
|
|
|
|
|
def centroid(site, verts):
|
|
area, c = 0, 0
|
|
v = verts - site
|
|
# print(v)
|
|
for i in range(len(v)):
|
|
x, y = v[i], v[(i + 1) % len(v)]
|
|
a = np.cross(x, y)
|
|
area += a
|
|
c += a * (x + y)
|
|
|
|
return c / (3 * area) + site
|
|
|
|
|
|
def flat_sheet(k):
|
|
x = np.linspace(-np.pi, np.pi, 101)
|
|
x = np.transpose([np.tile(a, len(a)), np.repeat(a, len(a))])
|
|
all_verts = np.hstack((a, np.zeros((len(a), 1))))
|
|
|
|
xx = np.array([[x for x in range(101 * 100) if (x - 100) % 101 > 0]]).T
|
|
|
|
b = np.hstack((xx, xx + 1, xx + 101))
|
|
c = np.hstack((xx + 1, xx + 102, xx + 101))
|
|
all_faces = np.concatenate((b, c))
|
|
|
|
return all_verts, all_faces
|
|
|
|
|
|
def torus_transform(v, R, r):
|
|
new_verts = np.empty(v.shape)
|
|
# Rotate by 90 degrees, and then shift to range [0, 2\pi]
|
|
v[:, :2] = np.matmul(np.array([[0, -1], [1, 0]]), v[:, :2].T).T
|
|
v[:, :2] = v[:, :2] % (2 * np.pi)
|
|
# v[:, 0], v[:, 1] = v[:, 1], v[:, 0]
|
|
|
|
new_verts[:, 0] = (R + r * np.cos(v[:, 1])) * np.cos(v[:, 0])
|
|
new_verts[:, 1] = (R + r * np.cos(v[:, 1])) * np.sin(v[:, 0])
|
|
new_verts[:, 2] = r * np.sin(v[:, 1])
|
|
return new_verts
|
|
|
|
|
|
def flat_region(c, verts, height):
|
|
m = len(verts)
|
|
|
|
centers = np.hstack((np.array([c, c]), np.array([[height], [0]])))
|
|
v_top = np.hstack((verts, height * np.ones((m, 1))))
|
|
v_bot = np.hstack((verts, np.zeros((m, 1))))
|
|
|
|
v = np.concatenate((centers, v_top, v_bot))
|
|
|
|
vi = np.atleast_2d(np.arange(m, dtype=int)).T
|
|
cent_top, cent_bot = np.zeros((m, 1), dtype=int), np.ones((m, 1), dtype=int)
|
|
vert_top_m, vert_top_p = vi + 2, (vi + 1) % m + 2
|
|
vert_bot_m, vert_bot_p = vi + m + 2, (vi + 1) % m + m + 2
|
|
|
|
top_face = np.hstack((cent_top, vert_top_m, vert_top_p))
|
|
bot_face = np.hstack((cent_bot, vert_bot_p, vert_bot_m))
|
|
|
|
f = np.concatenate((top_face, bot_face))
|
|
return v, f, 2
|
|
|
|
|
|
def torus_region(c, verts, height):
|
|
sd = SUBDIVISION_AMOUNT
|
|
d = verts - c
|
|
m = len(verts)
|
|
# 1 + 6 + 12 ... verts, 1 + 3 + 5 ... faces.
|
|
v, f = (
|
|
np.empty((1 + m * sd * (sd - 1) // 2, 2)),
|
|
np.empty((m * (sd - 1) ** 2, 3), dtype=int),
|
|
)
|
|
v[0] *= 0
|
|
|
|
# Vertices are ordered by rings going out from the center.
|
|
v_inds = np.arange(-1, sd)
|
|
v_inds = 1 + m * v_inds * (v_inds + 1) // 2
|
|
v_inds[0] = 0
|
|
|
|
for i in range(1, sd):
|
|
fj, fk = m * (i - 1) ** 2, m * i ** 2
|
|
region_verts = [
|
|
np.linspace(i * d[j] / (sd - 1), i * d[(j + 1) % m] / (sd - 1), i + 1)[:-1]
|
|
for j in range(m)
|
|
]
|
|
v[v_inds[i] : v_inds[i + 1]] = np.concatenate(region_verts)
|
|
|
|
faces = []
|
|
for j in range(m):
|
|
# v_off_inds = v_inds + j * np.arange(sd + 1)
|
|
r1, r2 = (
|
|
(j * (i - 1) + np.arange(i + 1)) % (m * (i - 1)) + v_inds[i - 1],
|
|
(j * i + np.arange(i + 2)) % (m * i) + v_inds[i],
|
|
)
|
|
r1, r2 = np.atleast_2d(r1).T, np.atleast_2d(r2).T
|
|
faces.append(np.hstack((r1[:-1], r2[:-2], r2[1:-1])))
|
|
faces.append(np.hstack((r1[:-2], r2[1:-2], r1[1:-1])))
|
|
|
|
f[fj:fk] = np.concatenate(faces)
|
|
|
|
v += c
|
|
v = np.hstack((v, height * np.ones((len(v), 1))))
|
|
return v, f, v_inds[sd - 1]
|
|
|
|
centers = np.array([[c[0], c[1], heights[i]]])
|
|
v_top = np.hstack((verts, heights[i] * np.ones((len(verts), 1))))
|
|
|
|
face_verts = np.concatenate((centers, v_top))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generates 3D model for equilibrium.")
|
|
parser.add_argument(
|
|
"sims_path", metavar="sim_dir", help="simulation to obtain equilibrium from"
|
|
)
|
|
parser.add_argument("frame_num", metavar="FRAME_NUM", type=int, help="frame number")
|
|
parser.add_argument("size", metavar="SIZE", type=int, help="width in mm")
|
|
parser.add_argument(
|
|
"--torus", dest="torus", action="store_true", help="generates torus model"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
torus = args.torus
|
|
size = args.size
|
|
|
|
# Get desired frame and load.
|
|
sim, frames = Simulation.load(args.sims_path)
|
|
frames = list(frames)
|
|
# num = 100
|
|
# nums = [838. 348. 14. 664, 725, 974]
|
|
# frames.sort(key=lambda x: x["energy"])
|
|
# frame = frames[-32]
|
|
# for i, f in enumerate(frames):
|
|
# if i == 858:
|
|
frame = frames[args.frame_num]
|
|
n, w, h = frame["domain"][0], frame["domain"][1], frame["domain"][2]
|
|
frame = sim.energy.mode(*frame["domain"], frame["arr"])
|
|
|
|
# Set up size and scaling variables.
|
|
dim = np.array([w, h])
|
|
vor = frame.vor_data
|
|
|
|
alpha = w / h
|
|
if torus:
|
|
r = alpha * size / 2
|
|
R = size / 2 - r
|
|
height_bounds = (0.5 * r, 1.25 * r)
|
|
scale = np.array([2 * np.pi, 2 * np.pi])
|
|
else:
|
|
height_bounds = (size / 10, 0.75 * size)
|
|
scale = np.array([alpha * size, size])
|
|
|
|
# Rescale and translate domain.
|
|
points = (vor.points - dim / 2) * scale / dim
|
|
vertices = (vor.vertices - dim / 2) * scale / dim
|
|
|
|
# Obtain height scaling.
|
|
heights = frame.stats["site_energies"] - ordered.e_hex(sim.domain)
|
|
heights -= np.min(heights)
|
|
heights *= (height_bounds[1] - height_bounds[0]) / np.max(heights)
|
|
heights += height_bounds[0]
|
|
|
|
# Prepare oriented regions and site_vert list
|
|
sites = points[:n]
|
|
regions = []
|
|
for i, x in enumerate(vor.point_region[:n]):
|
|
region = vor.regions[x]
|
|
x = sites[i]
|
|
p, q = vertices[region[0]] - x, vertices[region[1]] - x
|
|
regions.append(region if np.cross(p, q) >= 0 else region[::-1])
|
|
|
|
site_verts = []
|
|
for region in regions:
|
|
site_verts.append(vertices[region])
|
|
|
|
all_verts, all_faces = [], []
|
|
offsets = np.empty((n,), dtype=int)
|
|
for i, verts in enumerate(site_verts):
|
|
c = centroid(sites[i], verts)
|
|
if torus:
|
|
face_verts, faces, offset = torus_region(c, verts, heights[i])
|
|
else:
|
|
face_verts, faces, offset = flat_region(c, verts, heights[i])
|
|
|
|
old_len = len(all_verts)
|
|
if i == 0:
|
|
all_verts, all_faces = face_verts, faces
|
|
else:
|
|
all_verts = np.concatenate((all_verts, face_verts))
|
|
all_faces = np.concatenate((all_faces, faces + old_len))
|
|
|
|
offsets[i] = old_len + offset
|
|
|
|
# Merge regions
|
|
for i, x in enumerate(vor.ridge_points):
|
|
x, y = x[0], x[1] # Site indices.
|
|
no_x, no_y = x >= sim.domain.n, y >= sim.domain.n
|
|
if no_x and no_y:
|
|
continue
|
|
|
|
adj_verts = vor.ridge_vertices[i]
|
|
v1, v2 = vertices[adj_verts[0]], vertices[adj_verts[1]]
|
|
|
|
bot_face, top_face = None, None
|
|
if torus:
|
|
sd = SUBDIVISION_AMOUNT - 1
|
|
if no_x:
|
|
x, y = y, x
|
|
y = y % n
|
|
|
|
v1_x_ind = regions[x].index(adj_verts[0])
|
|
v2_x_ind = regions[x].index(adj_verts[1])
|
|
# Need to find matching vertices
|
|
v1, v2 = v1 % (2 * np.pi), v2 % (2 * np.pi)
|
|
y_verts = vertices[regions[y]] % (2 * np.pi)
|
|
v1_y_ind = np.argmin(np.linalg.norm(y_verts - v1, axis=1))
|
|
v2_y_ind = np.argmin(np.linalg.norm(y_verts - v2, axis=1))
|
|
|
|
# We want x lower than y.
|
|
off_x, off_y = offsets[x], offsets[y]
|
|
if heights[x] > heights[y]:
|
|
v1_x_ind, v1_y_ind = v1_y_ind, v1_x_ind
|
|
v2_x_ind, v2_y_ind = v2_y_ind, v2_x_ind
|
|
off_x, off_y = off_y, off_x
|
|
|
|
if not_in_order(v1_x_ind, v2_x_ind):
|
|
v1_x_ind, v2_x_ind = v2_x_ind, v1_x_ind
|
|
|
|
if not_in_order(v1_y_ind, v2_y_ind):
|
|
v1_y_ind, v2_y_ind = v2_y_ind, v1_y_ind
|
|
|
|
sds = np.atleast_2d(np.arange(sd)).T
|
|
xm, ym = sd * v1_x_ind + sds, sd * v1_y_ind + sds
|
|
xp, yp = xm + 1, ym + 1
|
|
xp[-1], yp[-1] = sd * v2_x_ind, sd * v2_y_ind
|
|
|
|
xm += off_x
|
|
xp += off_x
|
|
ym += off_y
|
|
yp += off_y
|
|
|
|
bot_face = np.hstack((xm, ym[::-1], xp))
|
|
top_face = np.hstack((xm, yp[::-1], ym[::-1]))
|
|
|
|
else:
|
|
if no_x ^ no_y:
|
|
if no_x:
|
|
x = y
|
|
m = len(regions[x])
|
|
v1_ind = regions[x].index(adj_verts[0])
|
|
v2_ind = regions[x].index(adj_verts[1])
|
|
if not_in_order(v1_ind, v2_ind):
|
|
v1_ind, v2_ind = v2_ind, v1_ind
|
|
|
|
v1_ind += offsets[x]
|
|
v2_ind += offsets[x]
|
|
bot_face = np.array([[v1_ind, m + v1_ind, m + v2_ind]], dtype=int)
|
|
top_face = np.array([[v1_ind, m + v2_ind, v2_ind]], dtype=int)
|
|
else:
|
|
v1_x_ind = regions[x].index(adj_verts[0])
|
|
v2_x_ind = regions[x].index(adj_verts[1])
|
|
v1_y_ind = regions[y].index(adj_verts[0])
|
|
v2_y_ind = regions[y].index(adj_verts[1])
|
|
|
|
# We want x lower than y.
|
|
off_x, off_y = offsets[x], offsets[y]
|
|
if heights[x] > heights[y]:
|
|
v1_x_ind, v1_y_ind = v1_y_ind, v1_x_ind
|
|
v2_x_ind, v2_y_ind = v2_y_ind, v2_x_ind
|
|
off_x, off_y = off_y, off_x
|
|
|
|
if not_in_order(v1_x_ind, v2_x_ind):
|
|
v1_x_ind, v2_x_ind = v2_x_ind, v1_x_ind
|
|
|
|
if not_in_order(v1_y_ind, v2_y_ind):
|
|
v1_y_ind, v2_y_ind = v2_y_ind, v1_y_ind
|
|
|
|
v1_x_ind += off_x
|
|
v2_x_ind += off_x
|
|
v1_y_ind += off_y
|
|
v2_y_ind += off_y
|
|
|
|
bot_face = np.array([[v1_x_ind, v1_y_ind, v2_x_ind]], dtype=int)
|
|
top_face = np.array([[v1_x_ind, v2_y_ind, v1_y_ind]], dtype=int)
|
|
|
|
all_faces = np.concatenate((all_faces, bot_face, top_face))
|
|
|
|
if torus:
|
|
all_verts = torus_transform(all_verts, R, all_verts[:, 2])
|
|
|
|
# Output into mesh
|
|
surf = mesh.Mesh(np.zeros(all_faces.shape[0], dtype=mesh.Mesh.dtype))
|
|
for i, f in enumerate(all_faces):
|
|
for j in range(3):
|
|
surf.vectors[i][j] = all_verts[f[j], :]
|
|
|
|
m_type = "Torus" if torus else "Flat"
|
|
out = OUTPUT_DIR / f"N{n} - {m_type} - {args.frame_num}.stl"
|
|
# temp = f"n{n}{m_type.lower()}{num}.stl"
|
|
surf.save(out)
|
|
# os.rename(temp, out)
|
|
print(f"Wrote to {out}.")
|
|
# n, alpha = 20, np.sqrt(3) * 2 / 5
|
|
# u, v = 5, 2
|
|
# domain = DomainParams(n, alpha, 1, 1)
|
|
# size = 1
|
|
|
|
# R, r = size, size * (1 - alpha) / alpha
|
|
|
|
# sites = ordered.sites(domain, (u, v))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
np.seterr(divide="ignore")
|
|
main()
|